Merge remote-tracking branch 'origin/main' into HEAD
This commit is contained in:
Yueqian Lin committed
commit
aef61597d3
1396 files changed
+151246
-39046
No files matched your search
Binary file not shown.
|
After Width: | Height: | Size: 1.4 MiB |
@@ -1,16 +0,0 @@
|
||||
[
|
||||
{
|
||||
"index": "vocal_240_51240",
|
||||
"language": "Mandarin",
|
||||
"time": [
|
||||
240,
|
||||
51240
|
||||
],
|
||||
"duration": "0.21 0.18 0.34 0.38 0.26 0.22 0.50 0.20 0.33 0.13 0.60 0.22 0.38 0.52 0.30 0.12 1.82 0.98 0.20 0.26 0.38 0.36 0.58 0.54 0.26 0.50 1.46 3.03 0.24 0.24 0.29 0.25 0.20 0.14 0.80 0.22 0.54 0.30 0.24 0.16 0.58 0.28 0.38 1.73 0.79 0.24 0.30 0.34 0.30 0.30 0.34 0.21 0.21 0.36 0.34 0.23 1.85 1.90 0.23 0.39 0.68 0.50 0.31 0.43 0.76 0.38 2.00 1.87 0.68 0.72 0.56 0.62 0.80 0.40 0.42 1.68 1.79 0.70 0.66 0.54 0.24 0.48 0.68 0.40 2.34 0.14",
|
||||
"text": "<SP> 只 是 因 为 为 在 人 群 群 中 多 看 了 你 一 眼 <SP> 再 也 没 能 忘 掉 你 容 颜 <SP> 梦 想 着 着 偶 偶 然 然 有 <SP> 一 一 天 再 相 见 <SP> 从 此 我 开 始 始 孤 孤 单 思 念 念 <SP> 想 想 你 时 你 你 在 天 边 <SP> 想 你 时 你 在 眼 前 前 <SP> 想 你 时 你 你 在 脑 海 <SP>",
|
||||
"phoneme": "<SP> zh_zhi3 zh_shi4 zh_yin1 zh_wei4 zh_wei2 zh_zai4 zh_ren2 zh_qun2 zh_qun2 zh_zhong1 zh_duo1 zh_kan4 zh_le5 zh_ni3 zh_yi1 zh_yan3 <SP> zh_zai4 zh_ye3 zh_mei2 zh_neng2 zh_wang4 zh_diao4 zh_ni3 zh_rong2 zh_yan2 <SP> zh_meng4 zh_xiang3 zh_zhe5 zh_zhe5 zh_ou3 zh_ou3 zh_ran2 zh_ran2 zh_you3 <SP> zh_yi1 zh_yi1 zh_tian1 zh_zai4 zh_xiang1 zh_jian4 <SP> zh_cong2 zh_ci3 zh_wo3 zh_kai1 zh_shi3 zh_shi3 zh_gu1 zh_gu1 zh_dan1 zh_si1 zh_nian4 zh_nian4 <SP> zh_xiang3 zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_ni3 zh_zai4 zh_tian1 zh_bian1 <SP> zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_zai4 zh_yan3 zh_qian2 zh_qian2 <SP> zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_ni3 zh_zai4 zh_nao3 zh_hai3 <SP>",
|
||||
"note_pitch": "0 64 64 64 66 68 66 66 66 64 64 64 64 66 64 60 61 0 63 63 63 64 66 63 61 59 56 0 68 66 66 68 68 66 66 64 64 0 64 66 61 61 61 64 0 62 63 63 64 64 66 65 66 61 58 59 56 0 69 71 66 68 68 71 66 64 61 0 66 61 68 66 64 63 61 59 0 71 66 68 68 71 66 64 61 0",
|
||||
"note_type": "1 2 2 2 2 3 2 2 2 3 2 2 2 2 2 2 2 1 2 2 2 2 2 2 2 2 2 1 2 2 2 3 2 3 2 3 2 1 2 3 2 2 2 2 1 2 2 2 2 2 3 2 3 2 2 2 3 1 2 3 2 2 2 3 2 2 2 1 2 2 2 2 2 2 2 3 1 2 2 2 2 3 2 2 2 1",
|
||||
"f0": "0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 318.9 332.2 331.2 323.9 312.3 0.0 0.0 0.0 0.0 0.0 0.0 340.9 342.3 340.2 337.4 333.5 329.3 327.8 328.3 329.1 330.5 331.9 331.5 329.2 327.6 327.4 327.7 329.6 330.5 329.5 328.1 328.5 330.1 330.4 329.5 328.9 331.0 333.6 332.8 331.7 330.2 330.6 330.9 329.5 327.6 332.4 343.5 363.9 368.7 369.3 368.2 366.7 371.3 382.0 392.9 402.3 409.8 413.8 416.3 416.4 415.4 414.1 413.4 415.3 417.4 418.3 413.6 416.0 0.0 0.0 0.0 0.0 368.2 362.5 361.3 362.0 362.9 364.3 366.4 367.5 368.4 368.9 368.1 367.6 369.7 371.1 371.1 369.5 368.0 367.3 366.6 367.1 369.5 372.0 371.3 368.2 368.1 369.6 369.9 373.4 376.5 360.2 285.6 0.0 0.0 0.0 0.0 0.0 0.0 0.0 384.0 376.7 372.0 371.7 373.9 376.3 375.5 361.2 333.0 317.8 319.0 329.1 333.5 333.2 333.2 333.3 328.0 314.5 0.0 0.0 0.0 320.6 326.5 332.9 335.7 334.1 329.9 327.4 326.3 326.6 329.1 330.5 330.0 328.0 326.7 328.6 331.1 332.4 332.1 332.8 333.2 333.1 331.8 328.6 318.9 293.2 323.6 327.5 330.0 332.8 331.1 319.4 272.2 0.0 0.0 0.0 0.0 278.5 294.5 310.0 318.5 322.7 327.4 334.1 337.2 336.3 328.5 321.0 324.2 341.5 362.2 374.7 375.8 370.5 366.0 365.2 368.7 371.7 374.4 373.0 368.7 370.4 374.8 375.1 372.5 368.7 363.0 357.1 359.0 366.8 377.1 379.5 371.5 359.9 351.0 358.4 371.7 377.9 375.8 367.8 359.7 363.6 375.3 379.9 377.5 372.9 359.8 345.2 337.1 334.5 333.4 332.8 329.6 319.2 292.8 261.7 253.4 262.8 273.0 278.1 279.3 278.6 277.6 277.9 278.3 278.0 277.3 276.7 275.6 275.3 276.1 277.8 278.2 277.7 277.6 277.2 276.4 275.5 273.9 271.8 270.7 274.7 282.7 285.5 281.4 273.1 264.4 262.2 268.9 278.4 284.3 283.6 277.3 268.5 262.2 263.2 270.7 278.8 285.0 287.3 282.0 272.1 265.2 266.0 272.1 278.2 283.8 285.1 281.4 272.6 267.2 269.1 276.1 282.7 289.3 290.5 285.6 273.9 267.5 267.4 271.8 277.7 282.6 283.1 277.8 269.8 265.6 268.6 273.8 281.1 286.2 287.2 282.3 270.4 264.5 263.3 263.0 268.5 268.3 268.7 270.9 277.2 277.4 284.3 285.7 282.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 311.5 306.4 306.0 306.8 305.8 305.0 307.2 310.8 312.4 313.5 314.4 314.6 313.7 312.2 311.1 311.1 312.8 314.6 312.7 306.5 299.7 302.7 308.7 315.2 314.2 310.9 309.4 309.8 310.6 310.7 310.8 311.4 312.2 313.0 312.5 312.2 312.2 311.9 312.0 312.2 311.4 306.1 303.3 309.3 315.0 325.3 324.4 323.9 325.0 326.4 327.7 327.7 328.1 328.6 327.8 328.3 328.3 327.6 328.6 336.1 347.2 359.4 371.0 374.9 371.3 366.5 369.5 375.5 376.2 370.2 365.7 367.3 372.9 376.9 375.3 368.4 360.1 358.0 365.4 380.8 383.5 380.4 374.4 367.6 365.9 372.0 376.3 377.9 375.7 372.6 368.6 361.2 353.7 0.0 309.4 302.3 299.2 300.6 305.7 308.5 308.6 309.0 310.9 310.9 309.5 308.6 309.0 311.3 313.0 314.4 314.5 313.0 312.1 311.2 308.0 299.5 295.8 295.0 285.7 278.0 280.1 281.2 279.9 278.3 278.3 279.5 279.7 272.2 259.1 0.0 0.0 0.0 240.7 236.0 233.8 233.5 234.6 235.7 235.6 236.2 239.8 245.4 248.8 250.1 250.4 249.9 249.1 247.2 241.7 231.7 216.0 197.6 192.1 197.2 204.3 206.5 205.4 201.7 200.8 203.5 208.2 210.9 210.3 207.1 203.5 202.3 202.1 202.8 205.1 207.3 208.1 207.3 205.4 202.1 199.3 200.5 206.1 212.9 212.9 206.5 197.1 191.1 191.2 196.8 203.6 207.8 208.7 206.4 201.7 197.3 197.4 200.7 206.5 209.9 209.9 206.9 203.3 201.2 203.4 207.7 210.5 208.9 207.6 206.5 204.2 202.0 203.6 205.7 210.8 213.1 214.3 210.6 204.1 199.7 202.3 211.9 217.7 215.1 215.0 215.3 213.3 0.0 0.0 216.3 0.0 0.0 195.2 209.0 205.3 201.3 196.8 195.7 195.8 195.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 431.9 425.8 421.0 417.7 419.3 422.7 423.4 421.1 417.9 411.9 408.6 411.5 414.9 419.6 426.7 424.2 0.0 0.0 0.0 280.8 279.4 369.7 371.8 370.6 369.8 370.5 370.5 371.4 374.6 374.0 361.0 332.7 328.0 0.0 359.0 365.0 371.1 373.8 371.9 369.7 369.6 371.0 378.3 390.4 403.4 414.1 417.9 417.5 417.0 416.6 415.7 414.9 414.5 413.8 413.5 413.6 413.8 415.2 416.9 417.3 415.4 415.0 414.1 411.4 409.2 403.5 392.7 375.6 367.5 365.8 368.8 370.7 370.8 369.0 367.0 366.0 366.2 367.8 369.2 368.8 367.2 366.7 367.3 367.3 368.8 368.5 366.8 365.3 364.0 363.0 365.0 367.3 367.7 365.5 364.0 364.9 367.4 370.3 371.3 369.9 366.5 364.1 368.3 385.8 406.7 409.4 399.7 376.0 357.2 359.9 372.4 377.5 366.9 345.1 320.6 315.8 321.8 329.0 329.9 326.4 324.5 324.8 324.9 325.3 325.1 324.7 325.9 326.8 326.6 326.6 326.4 325.4 323.2 322.5 323.9 326.2 328.6 329.9 329.5 329.0 328.2 327.7 327.9 329.6 331.2 331.4 334.2 340.1 335.5 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.Line truncated
|
||||
}
|
||||
]
|
||||
Binary file not shown.
@@ -1,16 +0,0 @@
|
||||
[
|
||||
{
|
||||
"index": "vocal_320_10687",
|
||||
"language": "Mandarin",
|
||||
"time": [
|
||||
320,
|
||||
10687
|
||||
],
|
||||
"duration": "0.23 0.34 0.26 0.70 0.52 0.46 0.36 0.44 0.14 0.24 0.64 0.47 0.51 1.10 0.28 0.38 0.32 0.32 0.38 0.32 0.31 0.19 1.45",
|
||||
"text": "<SP> 除 了 想 你 你 <SP> 除 了 了 爱 你 你 <SP> 我 什 么 什 么 都 愿 愿 意",
|
||||
"phoneme": "<SP> zh_chu2 zh_le5 zh_xiang3 zh_ni3 zh_ni3 <SP> zh_chu2 zh_le5 zh_le5 zh_ai4 zh_ni3 zh_ni3 <SP> zh_wo3 zh_shen2 zh_me5 zh_shen2 zh_me5 zh_dou1 zh_yuan4 zh_yuan4 zh_yi4",
|
||||
"note_pitch": "0 62 65 67 67 69 0 67 69 67 65 67 69 67 67 66 64 64 60 60 65 67 0",
|
||||
"note_type": "1 2 2 2 2 3 1 2 2 3 2 2 3 1 2 2 2 2 2 2 2 3 2",
|
||||
"f0": "0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 294.1 288.2 290.6 294.6 295.7 292.8 291.5 294.4 295.4 294.7 293.5 292.2 294.4 295.8 293.2 297.7 320.1 338.1 348.3 348.7 344.4 342.6 346.2 354.8 356.9 353.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 401.1 403.8 405.0 402.1 398.4 395.6 393.3 392.0 391.4 390.3 390.2 390.3 390.3 391.3 392.8 393.6 391.7 390.5 391.4 391.6 391.5 393.5 393.8 390.8 387.4 387.8 389.3 390.8 392.2 391.6 390.2 389.8 389.1 388.4 390.0 395.5 397.2 396.7 395.5 395.1 394.6 394.9 395.6 395.2 394.6 395.4 395.9 394.0 391.7 390.7 391.7 392.6 391.6 395.7 405.8 441.7 462.5 463.8 450.2 430.9 414.5 415.2 426.7 439.8 454.4 462.9 447.8 422.5 400.6 403.2 423.5 451.3 482.3 492.8 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 435.9 414.3 406.2 402.8 398.2 396.7 396.7 395.2 398.3 416.9 441.6 446.9 441.5 437.3 435.3 432.9 431.6 432.3 434.8 436.7 433.8 422.8 406.7 390.9 382.0 381.7 384.7 382.7 368.1 357.4 355.6 355.1 352.4 348.1 346.4 348.7 351.6 355.0 354.4 351.7 349.9 349.2 348.1 346.0 345.4 344.2 344.4 345.5 346.6 349.0 349.7 349.1 349.5 349.6 349.7 349.4 349.5 352.2 354.5 355.7 355.6 356.9 359.4 361.5 363.8 360.4 354.3 357.2 363.9 372.4 382.6 399.1 402.9 400.6 395.5 390.5 388.9 390.2 391.1 391.9 391.5 390.4 390.2 391.3 391.6 391.0 388.6 386.2 389.1 403.8 430.2 441.8 449.5 448.1 443.2 438.3 432.9 430.7 434.6 442.0 447.6 446.0 440.3 434.7 431.1 435.7 442.3 445.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 423.8 402.1 398.5 397.6 393.4 390.5 391.4 402.5 427.6 442.5 442.5 435.1 430.1 430.8 439.9 447.4 442.2 426.1 412.3 399.8 391.8 389.5 388.5 387.8 386.4 384.7 384.5 387.9 391.2 391.8 392.9 393.8 392.0 392.0 395.4 398.1 398.2 396.3 393.1 391.1 388.9 386.6 383.1 381.5 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 357.8 367.6 370.9 365.4 359.3 355.7 358.1 372.7 396.8 404.0 398.4 392.6 389.2 388.8 383.6 362.3 341.1 325.1 326.3 327.9 331.3 333.3 326.0 319.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 356.4 353.5 349.0 346.5 338.4 328.4 323.7 326.2 334.2 338.4 331.6 309.8 280.6 256.4 252.8 256.1 258.3 256.9 257.6 261.1 259.7 258.7 258.6 260.2 262.2 262.4 262.9 264.2 263.4 260.2 256.4 253.2 238.7 223.3 231.6 257.7 258.1 258.4 258.8 258.0 256.8 255.8 254.0 255.1 258.1 261.9 263.7 262.6 256.7 253.1 250.2 246.7 258.7 294.4 327.1 342.6 346.0 344.1 341.2 342.3 345.2 350.1 364.7 382.5 396.1 396.1 389.2 381.8 381.1 387.2 397.0 399.3 390.7 374.3 360.1 350.6 346.6 347.4 350.7 354.4 354.3 351.7 349.8 348.4 346.9 347.0 348.4 349.5 351.0 352.3 353.6 353.3 350.5 348.3 345.5 344.3 344.4 347.0 350.6 352.0 351.0 350.8 350.1 347.6 345.7 347.0 350.3 351.7 350.7 348.5 346.9 347.7 349.0 349.1 348.5 346.8 346.3 348.4 349.0 349.2 351.1 349.6 348.3 350.5 351.1 348.0 347.6 349.1 351.3 356.0 361.3 360.6 354.0 341.0 316.2 302.9 302.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0"
|
||||
}
|
||||
]
|
||||
Binary file not shown.
@@ -0,0 +1,211 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Benchmark legacy and automatic MP4 response encoding on synthetic input."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import statistics
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
from vllm_omni.entrypoints.openai.video_api_utils import (
|
||||
_encode_video_bytes,
|
||||
_encode_video_bytes_legacy,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BenchmarkVariant:
|
||||
label: str
|
||||
encoder: Callable[..., bytes]
|
||||
|
||||
|
||||
def _positive_int(value: str) -> int:
|
||||
parsed = int(value)
|
||||
if parsed < 1:
|
||||
raise argparse.ArgumentTypeError("value must be at least 1")
|
||||
return parsed
|
||||
|
||||
|
||||
def _build_inputs(
|
||||
*,
|
||||
frames: int,
|
||||
height: int,
|
||||
width: int,
|
||||
fps: int,
|
||||
audio_sample_rate: int,
|
||||
seed: int,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
rng = np.random.default_rng(seed)
|
||||
planar = rng.random((3, frames, height, width), dtype=np.float32)
|
||||
video = planar.transpose(1, 2, 3, 0)
|
||||
if not all(video[..., channel].flags.c_contiguous for channel in range(3)):
|
||||
raise RuntimeError("synthetic video does not expose contiguous channel planes")
|
||||
|
||||
sample_count = max(1, round(frames / fps * audio_sample_rate))
|
||||
timeline = np.arange(sample_count, dtype=np.float32) / audio_sample_rate
|
||||
left = np.sin(2 * np.pi * 220 * timeline).astype(np.float32, copy=False)
|
||||
right = np.sin(2 * np.pi * 330 * timeline).astype(np.float32, copy=False)
|
||||
return video, np.stack((left, right))
|
||||
|
||||
|
||||
def _measure(
|
||||
variant: BenchmarkVariant,
|
||||
*,
|
||||
video: np.ndarray,
|
||||
audio: np.ndarray,
|
||||
fps: int,
|
||||
audio_sample_rate: int,
|
||||
) -> tuple[bytes, dict[str, object]]:
|
||||
cpu_start = time.process_time_ns()
|
||||
wall_start = time.perf_counter_ns()
|
||||
output = variant.encoder(
|
||||
video,
|
||||
fps=fps,
|
||||
audio=audio,
|
||||
audio_sample_rate=audio_sample_rate,
|
||||
video_codec_options={"preset": "ultrafast", "threads": "0"},
|
||||
)
|
||||
wall_ms = (time.perf_counter_ns() - wall_start) / 1_000_000
|
||||
process_cpu_ms = (time.process_time_ns() - cpu_start) / 1_000_000
|
||||
return output, {
|
||||
"label": variant.label,
|
||||
"wall_ms": wall_ms,
|
||||
"process_cpu_ms": process_cpu_ms,
|
||||
"output_bytes": len(output),
|
||||
"output_sha256": hashlib.sha256(output).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 = [float(record["wall_ms"]) for record in selected]
|
||||
cpu_values = [float(record["process_cpu_ms"]) for record in selected]
|
||||
return {
|
||||
"runs": len(selected),
|
||||
"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("--frames", type=_positive_int, default=12)
|
||||
parser.add_argument("--height", type=_positive_int, default=96)
|
||||
parser.add_argument("--width", type=_positive_int, default=160)
|
||||
parser.add_argument("--fps", type=_positive_int, default=24)
|
||||
parser.add_argument("--audio-sample-rate", type=_positive_int, default=32000)
|
||||
parser.add_argument("--warmup", type=int, default=1)
|
||||
parser.add_argument("--rounds", type=_positive_int, default=3)
|
||||
parser.add_argument("--seed", type=int, default=20260817)
|
||||
parser.add_argument("--output-json", type=Path)
|
||||
args = parser.parse_args()
|
||||
if args.warmup < 0:
|
||||
parser.error("--warmup must be non-negative")
|
||||
return args
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = _parse_args()
|
||||
video, audio = _build_inputs(
|
||||
frames=args.frames,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
fps=args.fps,
|
||||
audio_sample_rate=args.audio_sample_rate,
|
||||
seed=args.seed,
|
||||
)
|
||||
baseline = BenchmarkVariant("legacy", _encode_video_bytes_legacy)
|
||||
candidate = BenchmarkVariant("automatic", _encode_video_bytes)
|
||||
|
||||
for _ in range(args.warmup):
|
||||
for variant in (baseline, candidate):
|
||||
_measure(
|
||||
variant,
|
||||
video=video,
|
||||
audio=audio,
|
||||
fps=args.fps,
|
||||
audio_sample_rate=args.audio_sample_rate,
|
||||
)
|
||||
|
||||
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 = _measure(
|
||||
variant,
|
||||
video=video,
|
||||
audio=audio,
|
||||
fps=args.fps,
|
||||
audio_sample_rate=args.audio_sample_rate,
|
||||
)
|
||||
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)}")
|
||||
|
||||
paired_savings_ms = []
|
||||
paired_savings_percent = []
|
||||
for round_index in range(args.rounds):
|
||||
pair = [record for record in records if record["round"] == round_index + 1]
|
||||
baseline_ms = float(next(record["wall_ms"] for record in pair if record["label"] == baseline.label))
|
||||
candidate_ms = float(next(record["wall_ms"] for record in pair if record["label"] == candidate.label))
|
||||
paired_savings_ms.append(baseline_ms - candidate_ms)
|
||||
paired_savings_percent.append((baseline_ms - candidate_ms) / baseline_ms * 100)
|
||||
|
||||
result = {
|
||||
"config": {
|
||||
"frames": args.frames,
|
||||
"height": args.height,
|
||||
"width": args.width,
|
||||
"fps": args.fps,
|
||||
"audio_sample_rate": args.audio_sample_rate,
|
||||
"warmup": args.warmup,
|
||||
"rounds": args.rounds,
|
||||
"seed": args.seed,
|
||||
"video_shape": list(video.shape),
|
||||
"video_strides": list(video.strides),
|
||||
"codec_options": {"preset": "ultrafast", "threads": "0"},
|
||||
},
|
||||
"records": records,
|
||||
"summary": {
|
||||
baseline.label: _summarize(records, baseline.label),
|
||||
candidate.label: _summarize(records, candidate.label),
|
||||
"paired_wall_saving_ms": {
|
||||
"median": statistics.median(paired_savings_ms),
|
||||
"min": min(paired_savings_ms),
|
||||
"max": max(paired_savings_ms),
|
||||
},
|
||||
"paired_wall_saving_percent": {
|
||||
"median": statistics.median(paired_savings_percent),
|
||||
"min": min(paired_savings_percent),
|
||||
"max": max(paired_savings_percent),
|
||||
},
|
||||
"output_sha256": output_hashes.pop(),
|
||||
},
|
||||
}
|
||||
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()
|
||||
@@ -0,0 +1,143 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.benchmarks.duplex import omni_duplex_eval_eval as eval_module
|
||||
from vllm_omni.benchmarks.duplex.omni_duplex_eval_clock import (
|
||||
extract_timed_sentences,
|
||||
normalize_response_items,
|
||||
split_text,
|
||||
validate_clock,
|
||||
)
|
||||
from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import (
|
||||
DuplexSample,
|
||||
canonical_task_type,
|
||||
family_for_split,
|
||||
load_samples,
|
||||
task_type_for_split,
|
||||
)
|
||||
from vllm_omni.benchmarks.duplex.omni_duplex_eval_metrics import (
|
||||
PROTOCOL_PIN,
|
||||
build_content_prompt,
|
||||
build_reminder_prompt,
|
||||
build_temporal_prompt,
|
||||
parse_judge_json,
|
||||
reminder_window,
|
||||
summarize_pr_results,
|
||||
temporal_window,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_dataset_mapping_and_manifest(tmp_path):
|
||||
manifest = tmp_path / "samples.json"
|
||||
manifest.write_text(json.dumps([{"id": "a", "split": "RTD_OCR", "video": "clip.mp4"}]), encoding="utf-8")
|
||||
sample = load_samples(manifest, media_root=tmp_path)[0]
|
||||
assert sample.family == "rtd"
|
||||
assert sample.video == str(tmp_path / "clip.mp4")
|
||||
assert family_for_split("PR_event_reminder") == "pr"
|
||||
assert canonical_task_type("post-event-reminder") == "post_event_reminder"
|
||||
assert task_type_for_split("PR_correction") == "correction"
|
||||
assert load_samples([{"id": "pr", "split": "PR_correction"}])[0].task_type == "correction"
|
||||
|
||||
|
||||
def test_hugging_face_subset_names_are_splits(monkeypatch):
|
||||
calls = []
|
||||
|
||||
def fake_load_dataset(name, **kwargs):
|
||||
calls.append((name, kwargs))
|
||||
return [{"id": "pr", "question_text": "Correct this."}]
|
||||
|
||||
monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
|
||||
sample = load_samples("Hothan/Omni-DuplexEval", split="PR_correction")[0]
|
||||
assert calls == [("Hothan/Omni-DuplexEval", {"split": "PR_correction"})]
|
||||
assert sample.split == "PR_correction"
|
||||
assert sample.task_type == "correction"
|
||||
|
||||
|
||||
def test_response_aliases_and_clock_guard():
|
||||
assert split_text("One. Two!") == ["One.", "Two!"]
|
||||
assert normalize_response_items({"chunks": [{"text": "x", "current_time": 800}]}) == [
|
||||
{"sentence": "x", "start": 800.0, "end": 800.0}
|
||||
]
|
||||
with pytest.raises(ValueError, match="clock=invalid"):
|
||||
validate_clock({"clock": "invalid"})
|
||||
validate_clock({"clock": "invalid"}, allow_invalid=True)
|
||||
timed = extract_timed_sentences([{"type": "response.output_text.delta", "delta": "Done.", "_media_clock_ms": 800}])
|
||||
assert [item.as_dict() for item in timed] == [{"sentence": "Done.", "start": 0.8, "end": 0.8}]
|
||||
with pytest.raises(ValueError, match="only clock=media"):
|
||||
extract_timed_sentences([], clock="wall")
|
||||
|
||||
|
||||
def test_protocol_windows_and_parsing():
|
||||
assert PROTOCOL_PIN == "ca3c122b4d4bf67afd6b18ea5e724b4561bdde48"
|
||||
assert temporal_window(4, 5, 10) == (2.0, 3.0)
|
||||
assert temporal_window(1, 1.2, 10) == (0.0, 0.5)
|
||||
assert reminder_window(5) == (5.0, 15.0)
|
||||
assert parse_judge_json('noise {"content_score": 2.345, "is_relevant": 1}') == {
|
||||
"content_score": 2.35,
|
||||
"is_relevant": 1,
|
||||
}
|
||||
assert parse_judge_json('noise "success_score": 1') == {"success_score": 1}
|
||||
assert summarize_pr_results([{"task_type": "correction", "all_success": 1}])["mean_all_success"] == 1.0
|
||||
|
||||
|
||||
def test_protocol_prompt_wording_is_pinned():
|
||||
prompts = [
|
||||
build_temporal_prompt(1, 2, "response", "question"),
|
||||
build_content_prompt("response", "question", ["one", "two"]),
|
||||
build_reminder_prompt("instruction", "response", "proactive_reminder"),
|
||||
build_reminder_prompt("instruction", "response", "post_event_reminder"),
|
||||
build_reminder_prompt("instruction", "response", "correction", "answer"),
|
||||
]
|
||||
assert [hashlib.sha256(prompt.encode()).hexdigest() for prompt in prompts] == [
|
||||
"5f7f5fcb4a37e4d2eea8ba23beabf74630e6fc2fb5b2661ceb94403254bb04d5",
|
||||
"1fccbed6c1ed8ea1c76d5564d7289ed6adca3678f6a0da2cc735942123d64457",
|
||||
"d0a0e3a9d873d2f368c5c99b9fe46681f683cc36a9fe0e09b99d6cbca1d6dd0c",
|
||||
"08f9951243c859d271f82ee53da22cd0bdab1fe93c6ca04f4870d3b959c0803a",
|
||||
"3de131a7e36cb4b8ac369bb05e87ede518be20221e214df3644211c179131a73",
|
||||
]
|
||||
|
||||
|
||||
def test_frame_sample_content_passes_frames(tmp_path, monkeypatch):
|
||||
response = tmp_path / "response.json"
|
||||
response.write_text(json.dumps([{"sentence": "A person moves.", "start": 1, "end": 2}]), encoding="utf-8")
|
||||
response.with_name("response.meta.json").write_text(json.dumps({"clock": "media"}), encoding="utf-8")
|
||||
video = tmp_path / "video.mp4"
|
||||
video.write_bytes(b"video")
|
||||
monkeypatch.setattr(eval_module, "extract_jpeg", lambda *args, **kwargs: b"\xff\xd8jpeg")
|
||||
|
||||
class Judge:
|
||||
def temporal(self, *args, **kwargs):
|
||||
return '{"temporal_score": 3, "is_relevant": 1}'
|
||||
|
||||
def content(self, prompt, video, frames=None, *, mode="video_url"):
|
||||
assert mode == "frame-sample"
|
||||
assert frames == [b"\xff\xd8jpeg", b"\xff\xd8jpeg"]
|
||||
return '{"content_score": 3}'
|
||||
|
||||
sample = DuplexSample(
|
||||
id="rtd",
|
||||
split="RTD_OCR",
|
||||
family="rtd",
|
||||
task_type=None,
|
||||
video=video,
|
||||
video_duration=4,
|
||||
)
|
||||
score = eval_module.evaluate_sample(
|
||||
sample,
|
||||
response,
|
||||
tmp_path / "score.json",
|
||||
Judge(),
|
||||
judge_video_mode="frame-sample",
|
||||
)
|
||||
assert score["content"]["frame_count"] == 2
|
||||
@@ -1,10 +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
|
||||
|
||||
"""
|
||||
Unit tests for metrics.py
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import pytest
|
||||
from vllm.benchmarks.serve import TaskType
|
||||
|
||||
@@ -66,6 +68,22 @@ def _make_output(prompt_len: int, output_tokens: int = 10) -> MixRequestFuncOutp
|
||||
return output
|
||||
|
||||
|
||||
def _calculate_test_metrics(outputs, goodput=None):
|
||||
return calculate_metrics(
|
||||
input_requests=[],
|
||||
outputs=outputs,
|
||||
dur_s=1.0,
|
||||
tokenizer=None,
|
||||
selected_percentiles=[50.0],
|
||||
goodput_config_dict=goodput or {},
|
||||
task_type=TaskType.GENERATION,
|
||||
selected_percentile_metrics=[],
|
||||
max_concurrency=None,
|
||||
request_rate=float("inf"),
|
||||
benchmark_duration=1.0,
|
||||
)[0]
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# total_input Tests
|
||||
# ============================================================================
|
||||
@@ -126,6 +144,228 @@ def test_audio_continuity_aggregation():
|
||||
assert p99 is not None and p99 > 0.4
|
||||
|
||||
|
||||
def test_unmeasured_duplex_latency_does_not_add_zero_samples():
|
||||
measured = _make_output(100, output_tokens=2)
|
||||
measured.ttft = 0.1
|
||||
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,
|
||||
}
|
||||
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,
|
||||
}
|
||||
|
||||
metrics = _calculate_test_metrics([measured, listen_only], {"ttft": 150.0})
|
||||
|
||||
assert metrics.mean_ttft_ms == 100.0
|
||||
assert metrics.mean_audio_ttfp_ms == 200.0
|
||||
assert metrics.mean_audio_rtf == 0.5
|
||||
assert (metrics.num_ttft_samples, metrics.num_audio_ttfp_samples, metrics.num_audio_rtf_samples) == (1, 1, 1)
|
||||
assert metrics.request_goodput == 1.0
|
||||
|
||||
|
||||
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,
|
||||
}
|
||||
|
||||
metrics = _calculate_test_metrics(
|
||||
[listen_only],
|
||||
{"ttft": float("inf"), "audio_ttft": float("inf")},
|
||||
)
|
||||
|
||||
assert (metrics.num_ttft_samples, metrics.num_audio_ttfp_samples, metrics.num_audio_rtf_samples) == (0, 0, 0)
|
||||
assert math.isnan(metrics.mean_ttft_ms)
|
||||
assert math.isnan(metrics.mean_audio_ttfp_ms)
|
||||
assert math.isnan(metrics.mean_audio_rtf)
|
||||
assert metrics.request_goodput == 0.0
|
||||
|
||||
|
||||
def test_unmeasured_duplex_tpot_does_not_add_zero_or_misalign_goodput():
|
||||
missing_tpot = _make_output(100, output_tokens=5)
|
||||
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}
|
||||
|
||||
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}
|
||||
|
||||
metrics = _calculate_test_metrics(
|
||||
[missing_tpot, slow_ttft],
|
||||
{"ttft": 500.0, "tpot": 200.0},
|
||||
)
|
||||
|
||||
assert metrics.num_tpot_samples == 1
|
||||
assert metrics.mean_tpot_ms == pytest.approx(100.0)
|
||||
assert metrics.request_goodput == 0.0
|
||||
|
||||
|
||||
def test_all_unmeasured_duplex_token_timing_is_not_reported_as_zero():
|
||||
output = _make_output(100, output_tokens=5)
|
||||
output.itl = []
|
||||
output.text_latency = output.ttft
|
||||
output.tpot_measured = False
|
||||
output.duplex_session_metrics = {"mean_ttft_ms": 100.0}
|
||||
|
||||
metrics = _calculate_test_metrics([output], {"tpot": float("inf")})
|
||||
|
||||
assert (metrics.num_tpot_samples, metrics.num_itl_samples) == (0, 0)
|
||||
assert math.isnan(metrics.mean_tpot_ms)
|
||||
assert math.isnan(metrics.mean_itl_ms)
|
||||
assert metrics.request_goodput == 0.0
|
||||
|
||||
|
||||
def test_duplex_response_timings_do_not_build_a_session_token_timeline():
|
||||
output = _make_output(100, output_tokens=5)
|
||||
output.latency = 101.0
|
||||
output.duplex_request_metrics = [
|
||||
{"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}
|
||||
|
||||
metrics = _calculate_test_metrics([output])
|
||||
|
||||
assert math.isnan(metrics.max_output_tokens_per_s)
|
||||
assert metrics.max_concurrent_requests == 1
|
||||
assert metrics.mean_tpot_ms == pytest.approx(100.0)
|
||||
assert metrics.mean_itl_ms == pytest.approx(100.0)
|
||||
|
||||
|
||||
def test_unmeasured_tpot_stays_missing_after_tokenizer_fallback():
|
||||
output = _make_output(100, output_tokens=0)
|
||||
output.generated_text = "timing metadata missing"
|
||||
output.itl = []
|
||||
output.text_latency = output.ttft
|
||||
output.tpot_measured = False
|
||||
output.duplex_session_metrics = {"mean_ttft_ms": 100.0}
|
||||
|
||||
def tokenizer(text, *, add_special_tokens):
|
||||
assert text == output.generated_text
|
||||
assert add_special_tokens is False
|
||||
return type("Tokenized", (), {"input_ids": [1, 2, 3]})()
|
||||
|
||||
metrics, _ = calculate_metrics(
|
||||
input_requests=[],
|
||||
outputs=[output],
|
||||
dur_s=1.0,
|
||||
tokenizer=tokenizer,
|
||||
selected_percentiles=[50.0],
|
||||
goodput_config_dict={"tpot": float("inf")},
|
||||
task_type=TaskType.GENERATION,
|
||||
selected_percentile_metrics=[],
|
||||
max_concurrency=None,
|
||||
request_rate=float("inf"),
|
||||
benchmark_duration=1.0,
|
||||
)
|
||||
|
||||
assert metrics.total_output == 3
|
||||
assert metrics.num_tpot_samples == 0
|
||||
assert math.isnan(metrics.mean_tpot_ms)
|
||||
assert metrics.request_goodput == 0.0
|
||||
|
||||
|
||||
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}
|
||||
|
||||
metrics = _calculate_test_metrics([output], {"tpot": 1.0})
|
||||
|
||||
assert (metrics.num_tpot_samples, metrics.num_itl_samples) == (0, 2)
|
||||
assert math.isnan(metrics.mean_tpot_ms)
|
||||
assert metrics.mean_itl_ms == 0.0
|
||||
assert metrics.request_goodput == 0.0
|
||||
|
||||
|
||||
def test_stage_output_tokens_without_client_timing_omit_tpot(capsys):
|
||||
output = _make_output(100, output_tokens=5)
|
||||
output.itl = []
|
||||
output.text_latency = 0.0
|
||||
output.ttft = 0.1
|
||||
|
||||
metrics, _ = calculate_metrics(
|
||||
input_requests=[],
|
||||
outputs=[output],
|
||||
dur_s=1.0,
|
||||
tokenizer=None,
|
||||
selected_percentiles=[50.0, 99.0],
|
||||
goodput_config_dict={},
|
||||
task_type=TaskType.GENERATION,
|
||||
selected_percentile_metrics=["tpot"],
|
||||
max_concurrency=None,
|
||||
request_rate=float("inf"),
|
||||
benchmark_duration=1.0,
|
||||
)
|
||||
|
||||
printed = capsys.readouterr().out
|
||||
assert metrics.num_tpot_samples == 0
|
||||
assert math.isnan(metrics.mean_tpot_ms)
|
||||
assert "Time per Output Token" not in printed
|
||||
assert "Mean TPOT" not in printed
|
||||
|
||||
|
||||
def test_consistent_client_latency_can_supply_tpot_fallback():
|
||||
output = _make_output(100, output_tokens=10)
|
||||
output.itl = []
|
||||
output.ttft = 0.1
|
||||
output.text_latency = 1.0
|
||||
|
||||
metrics = _calculate_test_metrics([output])
|
||||
|
||||
assert metrics.num_tpot_samples == 1
|
||||
assert metrics.mean_tpot_ms == pytest.approx(100.0)
|
||||
|
||||
|
||||
def test_single_token_responses_do_not_report_zero_tpot(capsys):
|
||||
output = _make_output(100, output_tokens=1)
|
||||
|
||||
metrics, _ = calculate_metrics(
|
||||
input_requests=[],
|
||||
outputs=[output],
|
||||
dur_s=1.0,
|
||||
tokenizer=None,
|
||||
selected_percentiles=[50.0],
|
||||
goodput_config_dict={},
|
||||
task_type=TaskType.GENERATION,
|
||||
selected_percentile_metrics=["tpot"],
|
||||
max_concurrency=None,
|
||||
request_rate=float("inf"),
|
||||
benchmark_duration=1.0,
|
||||
)
|
||||
|
||||
printed = capsys.readouterr().out
|
||||
assert metrics.num_tpot_samples == 0
|
||||
assert math.isnan(metrics.mean_tpot_ms)
|
||||
assert "Time per Output Token" not in printed
|
||||
|
||||
|
||||
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}
|
||||
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}
|
||||
|
||||
metrics = _calculate_test_metrics([text_only, audio_only], {"ttft": 500.0, "audio_ttft": 500.0})
|
||||
|
||||
assert metrics.request_goodput == 0.0
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# TTFT suppression for pure-audio (TTS) benchmarks
|
||||
# ============================================================================
|
||||
|
||||
@@ -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 patch.py
|
||||
@@ -17,6 +17,7 @@ from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput
|
||||
|
||||
from vllm_omni.benchmarks.patch.patch import (
|
||||
MixRequestFuncOutput,
|
||||
_apply_stage0_token_timings,
|
||||
async_request_openai_chat_omni_completions,
|
||||
async_request_openai_realtime_duplex,
|
||||
)
|
||||
@@ -102,7 +103,24 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
|
||||
received_at_s=now + 0.03,
|
||||
)
|
||||
self.events.add(
|
||||
{"type": "response.done", "response": {"id": response_id}},
|
||||
{
|
||||
"type": "response.done",
|
||||
"response": {
|
||||
"id": response_id,
|
||||
"metadata": {
|
||||
"vllm_omni": {
|
||||
"stage_metrics": {
|
||||
"0": {
|
||||
"num_tokens_out": 3,
|
||||
"vllm_ttft_ms": 20.0,
|
||||
"vllm_tpot_ms": 10.0,
|
||||
"vllm_itls_ms": [10.0, 10.0],
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
},
|
||||
},
|
||||
received_at_s=now + 0.04,
|
||||
)
|
||||
|
||||
@@ -143,6 +161,7 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
|
||||
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
|
||||
@@ -167,6 +186,10 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
|
||||
assert output.audio_ttfp > output.ttft
|
||||
assert output.audio_rtf > 0
|
||||
assert output.latency > 0
|
||||
assert output.output_tokens == 12
|
||||
assert output.itl == [0.01, 0.01] * 4
|
||||
assert output.text_latency == pytest.approx(output.ttft + 0.08)
|
||||
assert output.tpot_measured is True
|
||||
assert output.tts_turn_pcm_bytes == [b"\x00\x00" * 2400] * 4
|
||||
assert output.tts_output_pcm_bytes == b"\x00\x00" * 9600
|
||||
session_id = output.duplex_request_metrics[0]["session_id"]
|
||||
@@ -198,6 +221,54 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
|
||||
}
|
||||
|
||||
|
||||
def test_stage0_token_timings_use_weighted_tpot_when_itls_are_incomplete():
|
||||
output = MixRequestFuncOutput(ttft=0.1)
|
||||
|
||||
measured = _apply_stage0_token_timings(
|
||||
output,
|
||||
[
|
||||
{"output_token_count": 3, "itls_ms": [], "tpot_ms": 10.0},
|
||||
{"output_token_count": 2, "itls_ms": [], "tpot_ms": 40.0},
|
||||
],
|
||||
expected_output_tokens=5,
|
||||
)
|
||||
|
||||
assert measured is True
|
||||
assert output.itl == []
|
||||
assert output.tpot_measured is True
|
||||
assert (output.text_latency - output.ttft) / 4 == pytest.approx(0.02)
|
||||
|
||||
|
||||
def test_stage0_zero_itls_fall_back_to_positive_tpot():
|
||||
output = MixRequestFuncOutput(ttft=0.1)
|
||||
|
||||
measured = _apply_stage0_token_timings(
|
||||
output,
|
||||
[{"output_token_count": 3, "itls_ms": [0.0, 0.0], "tpot_ms": 25.0}],
|
||||
expected_output_tokens=3,
|
||||
)
|
||||
|
||||
assert measured is True
|
||||
assert output.itl == []
|
||||
assert output.tpot_measured is True
|
||||
assert (output.text_latency - output.ttft) / 2 == pytest.approx(0.025)
|
||||
|
||||
|
||||
def test_stage0_token_count_without_timing_is_not_measured():
|
||||
output = MixRequestFuncOutput(ttft=0.1)
|
||||
|
||||
measured = _apply_stage0_token_timings(
|
||||
output,
|
||||
[{"output_token_count": 3, "itls_ms": [], "tpot_ms": None}],
|
||||
expected_output_tokens=3,
|
||||
)
|
||||
|
||||
assert measured is False
|
||||
assert output.itl == []
|
||||
assert output.text_latency == output.ttft
|
||||
assert output.tpot_measured is False
|
||||
|
||||
|
||||
def create_sse_chunk(data_dict):
|
||||
"""Helper to create SSE formatted chunk"""
|
||||
return f"data: {json.dumps(data_dict)}\n\n".encode()
|
||||
|
||||
@@ -9,13 +9,17 @@ from pathlib import Path
|
||||
import pytest
|
||||
|
||||
from vllm_omni.benchmarks.data_modules.daily_omni_dataset import (
|
||||
DailyOmniDataset,
|
||||
daily_omni_local_qa_json,
|
||||
daily_omni_local_videos_dir,
|
||||
ensure_daily_omni_hub_root,
|
||||
resolve_daily_omni_local_root,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
|
||||
|
||||
_HUB_ID = "liarliar/Daily-Omni"
|
||||
|
||||
|
||||
def _write_qa_json(root: Path) -> Path:
|
||||
qa = root / "qa.json"
|
||||
@@ -23,6 +27,28 @@ def _write_qa_json(root: Path) -> Path:
|
||||
return qa
|
||||
|
||||
|
||||
def _stage_offline_hub_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, *, with_qa: bool = True) -> Path:
|
||||
"""Build a hub cache holding ``liarliar/Daily-Omni`` and cut off all network access.
|
||||
|
||||
Reproduces the air-gapped NPU runner: the dataset is pre-staged with ``hf download``, so it
|
||||
only exists in the *hub* cache -- never in the ``datasets`` cache that ``load_dataset`` reads.
|
||||
"""
|
||||
constants = pytest.importorskip("huggingface_hub").constants
|
||||
|
||||
cache = tmp_path / "hub"
|
||||
repo_dir = cache / "datasets--liarliar--Daily-Omni"
|
||||
snapshot = repo_dir / "snapshots" / "deadbeef"
|
||||
snapshot.mkdir(parents=True)
|
||||
if with_qa:
|
||||
_write_qa_json(snapshot)
|
||||
(repo_dir / "refs").mkdir()
|
||||
(repo_dir / "refs" / "main").write_text("deadbeef", encoding="utf-8")
|
||||
|
||||
monkeypatch.setattr(constants, "HF_HUB_CACHE", str(cache))
|
||||
monkeypatch.setattr(constants, "HF_HUB_OFFLINE", True)
|
||||
return snapshot
|
||||
|
||||
|
||||
def test_resolve_local_root_returns_none_for_hub_id() -> None:
|
||||
assert resolve_daily_omni_local_root("liarliar/Daily-Omni") is None
|
||||
assert resolve_daily_omni_local_root("") is None
|
||||
@@ -90,3 +116,37 @@ def test_missing_qa_json_reports_none(tmp_path: Path) -> None:
|
||||
root = resolve_daily_omni_local_root(str(tmp_path))
|
||||
assert daily_omni_local_qa_json(root) is None
|
||||
assert daily_omni_local_videos_dir(root) is None
|
||||
|
||||
|
||||
def test_hub_id_resolves_to_offline_cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
snapshot = _stage_offline_hub_cache(tmp_path, monkeypatch)
|
||||
|
||||
root = ensure_daily_omni_hub_root(_HUB_ID)
|
||||
|
||||
assert root == snapshot.resolve()
|
||||
assert daily_omni_local_qa_json(root) == (snapshot / "qa.json").resolve()
|
||||
|
||||
|
||||
def test_dataset_loads_hub_id_without_network(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""The regression this file exists for: a Hub id must not require ``load_dataset``."""
|
||||
snapshot = _stage_offline_hub_cache(tmp_path, monkeypatch)
|
||||
# Any use of `datasets` here means we went back to the path that cannot read the hub cache.
|
||||
monkeypatch.setattr("vllm_omni.benchmarks.data_modules.daily_omni_dataset.load_dataset", None)
|
||||
|
||||
dataset = DailyOmniDataset(dataset_path=_HUB_ID, random_seed=0)
|
||||
|
||||
assert dataset.qa_json_path == (snapshot / "qa.json").resolve()
|
||||
assert [row["video_id"] for row in dataset.data] == ["vid0"]
|
||||
|
||||
|
||||
def test_uncached_hub_id_reports_both_attempts(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_stage_offline_hub_cache(tmp_path, monkeypatch, with_qa=False)
|
||||
monkeypatch.setattr("vllm_omni.benchmarks.data_modules.daily_omni_dataset.load_dataset", None)
|
||||
|
||||
with pytest.raises(RuntimeError) as excinfo:
|
||||
DailyOmniDataset(dataset_path="someone-else/Daily-Omni", random_seed=0)
|
||||
|
||||
message = str(excinfo.value)
|
||||
assert "someone-else/Daily-Omni" in message
|
||||
assert "snapshot_download failed" in message
|
||||
assert "--daily-omni-qa-json" in message
|
||||
@@ -1,3 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Tests for Daily-Omni ``--daily-omni-pack-mode`` packing recipes.
|
||||
|
||||
The ``minicpm-interleave`` mode is an accuracy-critical protocol: MiniCPM-o only reaches
|
||||
@@ -27,6 +30,8 @@ _MODULE_NAME = "vllm_omni.benchmarks.data_modules.daily_omni_dataset"
|
||||
|
||||
if _MODULE_NAME not in sys.modules:
|
||||
_spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH)
|
||||
if _spec is None or _spec.loader is None:
|
||||
raise ImportError(f"Cannot load {_MODULE_NAME} from {_MODULE_PATH}")
|
||||
_mod = importlib.util.module_from_spec(_spec)
|
||||
sys.modules[_MODULE_NAME] = _mod
|
||||
_spec.loader.exec_module(_mod)
|
||||
@@ -87,7 +92,9 @@ def _stub_extract(monkeypatch, dataset: DailyOmniDataset, counter: list[int]) ->
|
||||
def fake_extract(video_path, *, audio_path, include_audio, **_kw):
|
||||
counter.append(1)
|
||||
frames = [Image.new("RGB", (16, 16), color=(i, i, i)) for i in range(_NUM_FRAMES)]
|
||||
segments = [np.zeros(1600, dtype=np.float32) for _ in range(_NUM_FRAMES)] if include_audio else []
|
||||
segments: list[np.ndarray] = (
|
||||
[np.zeros(1600, dtype=np.float32) for _ in range(_NUM_FRAMES)] if include_audio else []
|
||||
)
|
||||
return frames, segments
|
||||
|
||||
monkeypatch.setattr(dataset, "_extract_minicpm_frame_audio_segments", fake_extract)
|
||||
|
||||
@@ -0,0 +1,782 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import subprocess
|
||||
import tarfile
|
||||
import threading
|
||||
import time
|
||||
import wave
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput
|
||||
from vllm.benchmarks.serve import TaskType
|
||||
|
||||
from vllm_omni.benchmarks import omniinteract as oi
|
||||
from vllm_omni.benchmarks import serve as benchmark_serve
|
||||
from vllm_omni.benchmarks.data_modules import omniinteract_dataset as data
|
||||
from vllm_omni.benchmarks.metrics.metrics import calculate_metrics
|
||||
from vllm_omni.benchmarks.patch import patch as benchmark_patch
|
||||
from vllm_omni.entrypoints.cli.benchmark.cli_args import preprocess_serve_args
|
||||
from vllm_omni.entrypoints.cli.benchmark.serve import OmniBenchmarkServingSubcommand
|
||||
from vllm_omni.experimental.fullduplex.client import RealtimeEventCollector
|
||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
|
||||
MODEL = "openbmb/MiniCPM-o-4_5"
|
||||
|
||||
|
||||
def _case(tmp_path: Path, *, subset: str = "1q1a", name: str = "video.mp4") -> data.OmniInteractCase:
|
||||
video, annotation = tmp_path / name, tmp_path / f"{Path(name).stem}.json"
|
||||
video.parent.mkdir(parents=True, exist_ok=True)
|
||||
video.touch()
|
||||
annotation.write_text("{}")
|
||||
return data.OmniInteractCase(subset, name, video, annotation, "multi_turn")
|
||||
|
||||
|
||||
def _collector(*events: tuple[dict[str, object], float]) -> RealtimeEventCollector:
|
||||
collector = RealtimeEventCollector()
|
||||
for event, received_at in events:
|
||||
collector.add(event, received_at_s=received_at)
|
||||
return collector
|
||||
|
||||
|
||||
def _created(response_id: str = "r1") -> dict[str, object]:
|
||||
return {"type": "response.created", "response": {"id": response_id}}
|
||||
|
||||
|
||||
def _done(response_id: str = "r1", status: str = "completed") -> dict[str, object]:
|
||||
return {"type": "response.done", "response": {"id": response_id, "status": status}}
|
||||
|
||||
|
||||
def _audio(
|
||||
response_id: str = "r1",
|
||||
*,
|
||||
samples: int = 2400,
|
||||
rate: int | None = 24_000,
|
||||
value: int = 1,
|
||||
) -> dict[str, object]:
|
||||
event: dict[str, object] = {
|
||||
"type": "response.audio.delta",
|
||||
"response_id": response_id,
|
||||
"format": "pcm16",
|
||||
"delta": base64.b64encode(bytes((value, 0)) * samples).decode(),
|
||||
}
|
||||
if rate is not None:
|
||||
event["sample_rate_hz"] = rate
|
||||
return event
|
||||
|
||||
|
||||
def _text(response_id: str = "r1", value: str = "hello") -> dict[str, object]:
|
||||
return {"type": "response.audio_transcript.delta", "response_id": response_id, "delta": value}
|
||||
|
||||
|
||||
def _listen(*, buffering: bool = False) -> dict[str, object]:
|
||||
return {
|
||||
"type": "response.listen",
|
||||
"response": {"metadata": {"model_listen": True, "buffering": buffering}},
|
||||
}
|
||||
|
||||
|
||||
def _session_options(tmp_path: Path) -> data.OmniInteractSessionOptions:
|
||||
return data.OmniInteractSessionOptions(tmp_path / "artifacts", 10.0, 5.0, "reference.wav", True)
|
||||
|
||||
|
||||
def _sample(case: data.OmniInteractCase, options: data.OmniInteractSessionOptions):
|
||||
return data.OmniInteractSampleRequest(
|
||||
prompt="",
|
||||
prompt_len=0,
|
||||
expected_output_len=0,
|
||||
multi_modal_data=None,
|
||||
request_id="measured-0",
|
||||
omniinteract_case=case,
|
||||
omniinteract_options=options,
|
||||
)
|
||||
|
||||
|
||||
def _write_dataset(root: Path) -> Path:
|
||||
data_root = root / "data"
|
||||
for subset in ("1q1a", "1q1a_math"):
|
||||
subset_root = data_root / subset
|
||||
(subset_root / "videos").mkdir(parents=True)
|
||||
(subset_root / "annotations").mkdir()
|
||||
(subset_root / "videos" / f"{subset}.mp4").touch()
|
||||
(subset_root / "annotations" / f"{subset}.json").write_text("{}")
|
||||
(subset_root / "video_json_map.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"entries": [
|
||||
{
|
||||
"video": f"videos/{subset}.mp4",
|
||||
"annotation": f"annotations/{subset}.json",
|
||||
"scene_type": "multi_turn",
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
)
|
||||
one_to_many = data_root / "1qna"
|
||||
(one_to_many / "videos_bench" / "nested").mkdir(parents=True)
|
||||
(one_to_many / "annotations" / "nested").mkdir(parents=True)
|
||||
(one_to_many / "videos_bench" / "nested" / "guide.mp4").touch()
|
||||
(one_to_many / "annotations" / "nested" / "guide.json").write_text("{}")
|
||||
return data_root
|
||||
|
||||
|
||||
def _successful_output(
|
||||
tmp_path: Path,
|
||||
*,
|
||||
audio_time: float = 10.0,
|
||||
status: str = "completed",
|
||||
) -> tuple[data.OmniInteractCase, RealtimeEventCollector, oi.OmniInteractCaseResult]:
|
||||
case = _case(tmp_path)
|
||||
collector = _collector(
|
||||
(_created(), 9.9),
|
||||
(_audio(), audio_time),
|
||||
(_text(), audio_time + 0.01),
|
||||
(_done(status=status), audio_time + 0.1),
|
||||
)
|
||||
result = oi.OmniInteractCaseResult(case.subset, str(case.video_path), "")
|
||||
return case, collector, result
|
||||
|
||||
|
||||
def _write_success(
|
||||
root: Path,
|
||||
case: data.OmniInteractCase,
|
||||
collector: RealtimeEventCollector,
|
||||
result: oi.OmniInteractCaseResult,
|
||||
**kwargs,
|
||||
):
|
||||
require_response = kwargs.pop("require_response", True)
|
||||
return oi.prepare_success_artifacts(
|
||||
root,
|
||||
case,
|
||||
collector,
|
||||
stream_start=10.0,
|
||||
video_duration_s=1.0,
|
||||
require_response=require_response,
|
||||
result=result,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def test_dataset_discovers_official_layouts_and_total_selection(tmp_path: Path):
|
||||
root = _write_dataset(tmp_path)
|
||||
all_cases = data.discover_omniinteract_cases(root, data.OMNIINTERACT_SUBSETS, num_prompts=0, disable_shuffle=True)
|
||||
selected = data.discover_omniinteract_cases(root, data.OMNIINTERACT_SUBSETS, num_prompts=2, disable_shuffle=True)
|
||||
assert [case.subset for case in all_cases] == ["1q1a", "1q1a_math", "1qna"]
|
||||
assert all_cases[-1].video_rel == "videos_bench/nested/guide.mp4"
|
||||
assert [case.subset for case in selected] == ["1q1a", "1q1a_math"]
|
||||
|
||||
|
||||
def _archive(path: Path, member: str) -> None:
|
||||
with tarfile.open(path, "w") as archive:
|
||||
content = b"{}"
|
||||
item = tarfile.TarInfo(member)
|
||||
item.size = len(content)
|
||||
archive.addfile(item, io.BytesIO(content))
|
||||
|
||||
|
||||
def test_archive_is_safe_and_atomically_shared(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
bad = tmp_path / "bad.tar"
|
||||
_archive(bad, "../escape")
|
||||
with pytest.raises(ValueError, match="Unsafe path"):
|
||||
data._extract_archive(bad, tmp_path / "bad-cache")
|
||||
|
||||
good = tmp_path / "good.tar"
|
||||
_archive(good, "data/1q1a/video_json_map.json")
|
||||
target = tmp_path / "cache"
|
||||
barrier, safe_extract = threading.Barrier(2), data._safe_extract
|
||||
monkeypatch.setattr(data, "_safe_extract", lambda *args: (barrier.wait(), safe_extract(*args)))
|
||||
with ThreadPoolExecutor(2) as executor:
|
||||
roots = list(executor.map(lambda _: data._extract_archive(good, target), range(2)))
|
||||
assert roots[0] == roots[1]
|
||||
assert (roots[0] / "1q1a" / "video_json_map.json").is_file()
|
||||
assert not list(target.glob(".tmp-*"))
|
||||
|
||||
|
||||
def test_hub_archive_uses_vllm_filesystem(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
source = tmp_path / "source.tar"
|
||||
_archive(source, "data/1q1a/video_json_map.json")
|
||||
|
||||
class FS:
|
||||
def get_file(self, remote: str, local: str) -> None:
|
||||
assert remote == "datasets/org/repo/data.tar.gz"
|
||||
Path(local).write_bytes(source.read_bytes())
|
||||
|
||||
monkeypatch.setenv("HF_HOME", str(tmp_path / "hf"))
|
||||
monkeypatch.setattr(data, "hf_fs", lambda: FS())
|
||||
assert (data.resolve_omniinteract_root(None, "org/repo") / "1q1a").is_dir()
|
||||
|
||||
|
||||
def test_media_commands_are_bounded(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
|
||||
calls: list[float] = []
|
||||
commands: list[list[str]] = []
|
||||
bounded_run = oi._run_media_command
|
||||
|
||||
def run(command, *, timeout_s, text=False):
|
||||
calls.append(timeout_s)
|
||||
commands.append(command)
|
||||
if command[0] == "ffprobe":
|
||||
return subprocess.CompletedProcess(command, 0, "1.0", "")
|
||||
return subprocess.CompletedProcess(command, 0, b"", b"")
|
||||
|
||||
monkeypatch.setattr(oi, "_run_media_command", run)
|
||||
duration, pcm, frames = oi.prepare_media(tmp_path / "video.mp4", 1.0, timeout_s=3.0, max_duration_s=1.0)
|
||||
assert (duration, len(pcm), frames, calls) == (1.0, 32_000, [None], [3.0, 3.0, 3.0])
|
||||
assert commands[1][commands[1].index("-t") + 1] == "1.0"
|
||||
assert commands[1][commands[1].index("-fs") + 1] == "32000"
|
||||
|
||||
monkeypatch.setattr(subprocess, "run", lambda *a, **k: (_ for _ in ()).throw(subprocess.TimeoutExpired("x", 3)))
|
||||
with pytest.raises(TimeoutError, match="timed out"):
|
||||
bounded_run(["ffmpeg"], timeout_s=3.0)
|
||||
|
||||
|
||||
def test_media_duration_is_bounded_before_decode(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
|
||||
calls = 0
|
||||
|
||||
def run(command, *, timeout_s, text=False):
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
return subprocess.CompletedProcess(command, 0, "601.0", "")
|
||||
|
||||
monkeypatch.setattr(oi, "_run_media_command", run)
|
||||
with pytest.raises(ValueError, match="Invalid video duration"):
|
||||
oi.prepare_media(tmp_path / "video.mp4", 1.0, timeout_s=3.0, max_duration_s=600.0)
|
||||
assert calls == 1
|
||||
|
||||
|
||||
def test_final_commit_keeps_a_partial_model_unit():
|
||||
"""Exact units retain a sample for the server's correlated final flush."""
|
||||
|
||||
events = [{"session": {"capabilities": {"chunk_period_ms": 1000}}}]
|
||||
exact_unit = b"\0\0" * 16_000
|
||||
assert oi._ensure_final_commit_tail(exact_unit, events) == exact_unit[:-2]
|
||||
assert oi._ensure_final_commit_tail(exact_unit + b"\0\0", events) == exact_unit + b"\0\0"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("events", "match"),
|
||||
[
|
||||
(((_done(), 0.0),), "without response.created"),
|
||||
(((_created(), 0.0), (_created(), 0.1)), "duplicate response.created"),
|
||||
(((_created(), 0.0), (_done(), 0.1), (_done(), 0.2)), "duplicate response.done"),
|
||||
(
|
||||
(
|
||||
(_created(), 0.0),
|
||||
({"type": "response.done", "response": {"id": "r1", "status_details": {"type": "failed"}}}, 0.1),
|
||||
),
|
||||
"reports failure",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_response_ledger_rejects_identity_errors(events, match: str):
|
||||
with pytest.raises(ValueError, match=match):
|
||||
oi.response_ledger(_collector(*events))
|
||||
|
||||
|
||||
class _CompletionClient:
|
||||
def __init__(self, collector: RealtimeEventCollector):
|
||||
self.events, self.acks = collector, []
|
||||
|
||||
def raise_if_reader_stopped(self) -> None:
|
||||
return None
|
||||
|
||||
async def send_playback_ack(self, response_id: str, played_ms: int) -> None:
|
||||
self.acks.append((response_id, played_ms))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("tail", "succeeds"),
|
||||
[
|
||||
(((_created(), 1.1), (_done(), 1.2)), True),
|
||||
(((_listen(), 1.1),), True),
|
||||
(((_listen(buffering=True), 1.1),), False),
|
||||
((({"type": "session.closed", "reason": "timeout"}, 1.1),), False),
|
||||
],
|
||||
)
|
||||
async def test_completion_requires_a_final_model_decision(tail, succeeds: bool):
|
||||
events = _collector(({"type": "input_audio_buffer.committed"}, 1.0), *tail)
|
||||
coroutine = oi.wait_for_session_completion(
|
||||
_CompletionClient(events), oi._Playback(), commit_from=0, timeout_s=0.08, settle_s=0.0
|
||||
)
|
||||
if succeeds:
|
||||
assert await coroutine == 0
|
||||
else:
|
||||
with pytest.raises((TimeoutError, RuntimeError)):
|
||||
await coroutine
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deferred_commit_ignores_the_prior_response_terminal():
|
||||
events = _collector(
|
||||
(_created("old"), 0.0),
|
||||
({"type": "input_audio_buffer.committed", "event": {"overlap_deferred": True}}, 1.0),
|
||||
(_done("old"), 1.1),
|
||||
(_listen(), 1.2),
|
||||
)
|
||||
assert (
|
||||
await oi.wait_for_session_completion(
|
||||
_CompletionClient(events), oi._Playback(), commit_from=1, timeout_s=0.1, settle_s=0.0
|
||||
)
|
||||
== 1
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_playback_acks_only_after_serial_audio_drain():
|
||||
collector = _collector((_audio(samples=2400), 10.0), (_done(), 10.01))
|
||||
client, playback = _CompletionClient(collector), oi._Playback()
|
||||
await playback.acknowledge(client, now=10.05)
|
||||
assert client.acks == []
|
||||
await playback.acknowledge(client, now=10.1)
|
||||
assert client.acks == [("r1", 100)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("event", "match"),
|
||||
[
|
||||
(_audio(rate=16_000), "24000 Hz"),
|
||||
({**_audio(), "format": "opus"}, "pcm16"),
|
||||
({**_audio(), "delta": "!"}, "base64"),
|
||||
],
|
||||
)
|
||||
def test_playback_rejects_non_official_audio(event: dict[str, object], match: str):
|
||||
playback = oi._Playback()
|
||||
with pytest.raises(ValueError, match=match):
|
||||
playback.ingest(_collector((event, 1.0)))
|
||||
|
||||
|
||||
def test_artifacts_publish_official_bundle_and_sparse_deferred_state(tmp_path: Path):
|
||||
case, collector, result = _successful_output(tmp_path)
|
||||
summary = _write_success(
|
||||
tmp_path / "out",
|
||||
case,
|
||||
collector,
|
||||
result,
|
||||
)
|
||||
context = result._artifact_context
|
||||
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")
|
||||
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)
|
||||
with wave.open(str(directory / "output.wav")) as output:
|
||||
assert (output.getframerate(), output.getnframes()) == (24_000, 24_000)
|
||||
transcript = json.loads((directory / "wav_transcript.json").read_text())
|
||||
assert transcript["text"] == "hello" and transcript["chunks"][0]["timestamp"] == [0.0, 0.1]
|
||||
assert "serialized playback queue" in transcript["timestamp_semantics"]
|
||||
assert result._artifact_context is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("audio_time", "status", "reason"),
|
||||
[(10.95, "completed", "audio_clipped"), (10.0, "cancelled", "cancelled_response")],
|
||||
)
|
||||
def test_ineligible_outputs_are_excluded_from_manifest(tmp_path: Path, audio_time: float, status: str, reason: str):
|
||||
case, collector, result = _successful_output(tmp_path, audio_time=audio_time, status=status)
|
||||
_write_success(tmp_path / "out", case, collector, result, require_response=False)
|
||||
assert reason in result.official_eval_ineligible_reasons
|
||||
oi.write_batch_artifacts(tmp_path / "out", [case], [result])
|
||||
assert (tmp_path / "out" / "official_eval_manifest.jsonl").read_text() == ""
|
||||
|
||||
|
||||
def test_artifacts_require_complete_audio_and_transcript_response(tmp_path: Path):
|
||||
case, _, result = _successful_output(tmp_path)
|
||||
collector = _collector((_created(), 10.0), (_audio(), 10.0), (_done(), 10.1))
|
||||
with pytest.raises(ValueError, match="audio and transcript"):
|
||||
_write_success(tmp_path, case, collector, result)
|
||||
|
||||
|
||||
def test_publication_failure_removes_partial_success_bundle(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
case, collector, result = _successful_output(tmp_path)
|
||||
_write_success(tmp_path, case, collector, result)
|
||||
original = oi._atomic_write_json
|
||||
|
||||
def fail(path: Path, value: object) -> None:
|
||||
if path.name == "events.json":
|
||||
raise OSError("disk full")
|
||||
original(path, value)
|
||||
|
||||
monkeypatch.setattr(oi, "_atomic_write_json", fail)
|
||||
with pytest.raises(OSError, match="disk full"):
|
||||
oi.publish_deferred_case_artifacts(tmp_path, case, result)
|
||||
directory = oi._output_dir(tmp_path, case)
|
||||
assert not any((directory / name).exists() for name in oi.SUCCESS_ARTIFACTS)
|
||||
assert not result.success and "artifact_write_failed" in result.official_eval_ineligible_reasons
|
||||
|
||||
|
||||
def test_atomic_write_preserves_destination_on_failure(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
destination = tmp_path / "output.wav"
|
||||
destination.write_bytes(b"old")
|
||||
monkeypatch.setattr(oi, "write_pcm16_wav", lambda *a, **k: (_ for _ in ()).throw(OSError("write")))
|
||||
with pytest.raises(OSError):
|
||||
oi._atomic_write_wav(destination, b"\0\0", 24_000)
|
||||
assert destination.read_bytes() == b"old"
|
||||
assert not list(tmp_path.glob(".*.tmp"))
|
||||
|
||||
|
||||
class _RealtimeClient:
|
||||
instances: list[_RealtimeClient] = []
|
||||
|
||||
def __init__(self, url: str, **kwargs):
|
||||
self.url, self.events, self.acks, self.configure_kwargs = url, RealtimeEventCollector(), [], {}
|
||||
self.instances.append(self)
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args):
|
||||
return None
|
||||
|
||||
async def configure(self, model: str, **kwargs) -> None:
|
||||
self.configure_kwargs = kwargs
|
||||
self.events.add({"type": "session.created", "session": {"capabilities": {"chunk_period_ms": 1000}}})
|
||||
|
||||
async def send(self, event: dict[str, object]) -> None:
|
||||
return None
|
||||
|
||||
async def commit(self) -> None:
|
||||
now = time.monotonic()
|
||||
for event, offset in (
|
||||
({"type": "input_audio_buffer.committed"}, 0.0),
|
||||
(_created(), 0.001),
|
||||
(_audio(samples=24), 0.002),
|
||||
(_text(value="answer"), 0.003),
|
||||
(_done(), 0.004),
|
||||
):
|
||||
self.events.add(event, received_at_s=now + offset)
|
||||
|
||||
async def send_playback_ack(self, response_id: str, played_ms: int) -> None:
|
||||
self.acks.append((response_id, played_ms))
|
||||
|
||||
async def close_session(self, **kwargs) -> None:
|
||||
self.events.add({"type": "session.closed"})
|
||||
|
||||
def raise_if_reader_stopped(self) -> None:
|
||||
return None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_public_runner_executes_one_prepared_session(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
case = _case(tmp_path)
|
||||
config = oi.OmniInteractBenchmarkConfig(
|
||||
model=MODEL,
|
||||
output_root=tmp_path / "out",
|
||||
ref_audio="ref.wav",
|
||||
require_response=True,
|
||||
extra_body={"custom": "value"},
|
||||
)
|
||||
prepared = data.OmniInteractPreparedInput(1.0, b"\0\0\0\0", ("frame",), "data:audio/wav;base64,ref")
|
||||
monkeypatch.setattr(oi, "RealtimeDuplexClient", _RealtimeClient)
|
||||
monkeypatch.setattr(oi, "_COMPLETION_SETTLE_S", 0.0)
|
||||
result = await oi.run_omniinteract_case(
|
||||
case,
|
||||
config,
|
||||
request_index=0,
|
||||
prepared_input=prepared,
|
||||
)
|
||||
assert result.success and result.transcript == "answer" and result._artifact_context is not None
|
||||
assert not config.output_root.exists()
|
||||
assert "autostart=0" in _RealtimeClient.instances[-1].url
|
||||
assert _RealtimeClient.instances[-1].configure_kwargs["extra_body"] == {"custom": "value"}
|
||||
assert _RealtimeClient.instances[-1].acks == [("r1", 1)]
|
||||
|
||||
|
||||
def test_standard_sample_loading_prepares_media_before_timing(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
_write_dataset(tmp_path)
|
||||
ref = tmp_path / "ref.wav"
|
||||
ref.touch()
|
||||
parser = TrackingArgumentParser()
|
||||
OmniBenchmarkServingSubcommand.add_cli_args(parser)
|
||||
args = parser.parse_args(
|
||||
[
|
||||
"--backend",
|
||||
"openai-realtime-duplex",
|
||||
"--dataset-name",
|
||||
"omniinteract",
|
||||
"--dataset-path",
|
||||
str(tmp_path),
|
||||
"--model",
|
||||
MODEL,
|
||||
"--endpoint",
|
||||
"/v1/realtime",
|
||||
"--num-prompts",
|
||||
"0",
|
||||
"--omniinteract-ref-audio",
|
||||
str(ref),
|
||||
]
|
||||
)
|
||||
preprocess_serve_args(args)
|
||||
assert args.endpoint == "/v1/realtime"
|
||||
monkeypatch.setattr(benchmark_patch, "reference_audio_data_url", lambda _: "data:audio/wav;base64,ref")
|
||||
monkeypatch.setattr(benchmark_patch, "prepare_media", lambda *a, **k: (1.0, b"pcm", ["frame"]))
|
||||
samples = benchmark_patch.get_samples(args, None)
|
||||
assert len(samples) == args.num_prompts == 3
|
||||
assert all(sample.omniinteract_prepared_input.video_frames == ("frame",) for sample in samples)
|
||||
|
||||
|
||||
def _request(case: data.OmniInteractCase, options: data.OmniInteractSessionOptions) -> RequestFuncInput:
|
||||
request = RequestFuncInput(
|
||||
model=MODEL,
|
||||
model_name=MODEL,
|
||||
prompt="",
|
||||
api_url="http://server:8000/v1/realtime",
|
||||
prompt_len=0,
|
||||
output_len=0,
|
||||
logprobs=None,
|
||||
multi_modal_content=None,
|
||||
ignore_eos=False,
|
||||
request_id="measured-0",
|
||||
)
|
||||
request.omniinteract_case, request.omniinteract_options = case, options
|
||||
return request
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_adapter_requires_and_forwards_exact_prepared_payload(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
case, options = _case(tmp_path), _session_options(tmp_path)
|
||||
request = _request(case, options)
|
||||
output = await benchmark_patch.async_request_openai_realtime_duplex(request, session=None)
|
||||
assert not output.success and "prepared before benchmark request timing" in output.error
|
||||
|
||||
prepared = data.OmniInteractPreparedInput(1.0, b"pcm", ("frame",), "data:audio/wav;base64,ref")
|
||||
request.omniinteract_prepared_input = prepared
|
||||
request.extra_headers = {"X-Proxy": "test"}
|
||||
request.extra_body = {"custom": "value"}
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "secret")
|
||||
seen = {}
|
||||
|
||||
async def run(_case_arg, config, **kwargs):
|
||||
seen["headers"] = config.extra_headers
|
||||
seen["extra_body"] = config.extra_body
|
||||
seen.update(kwargs)
|
||||
return oi.OmniInteractCaseResult(
|
||||
case.subset,
|
||||
str(case.video_path),
|
||||
str(options.output_root),
|
||||
success=True,
|
||||
transcript="timing metadata missing",
|
||||
output_tokens=0,
|
||||
duplex_request_metrics=[{"request_metrics": {"ttft_ms": 1.0}}],
|
||||
duplex_session_metrics={"mean_ttft_ms": 10.0},
|
||||
)
|
||||
|
||||
monkeypatch.setattr(benchmark_patch, "run_omniinteract_case", run)
|
||||
output = await benchmark_patch.async_request_openai_realtime_duplex(request, session=None)
|
||||
assert output.success and seen["prepared_input"] is prepared
|
||||
assert seen["headers"] == {"Authorization": "Bearer secret", "X-Proxy": "test", "x-request-id": "measured-0"}
|
||||
assert seen["extra_body"] == {"custom": "value"}
|
||||
assert output.itl == []
|
||||
assert output.tpot_measured is False
|
||||
assert seen["capture_artifacts"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("request_metrics", "expected_itl", "expected_text_latency", "expected_tpot_measured"),
|
||||
[
|
||||
(
|
||||
[
|
||||
{"stage0_tokens": {"output_token_count": 3, "itls_ms": [10.0, 20.0]}},
|
||||
{"stage0_tokens": {"output_token_count": 2, "itls_ms": [30.0]}},
|
||||
],
|
||||
[0.01, 0.02, 0.03],
|
||||
0.16,
|
||||
True,
|
||||
),
|
||||
(
|
||||
[
|
||||
{"stage0_tokens": {"output_token_count": 3, "tpot_ms": 10.0}},
|
||||
{"stage0_tokens": {"output_token_count": 2, "tpot_ms": 40.0}},
|
||||
],
|
||||
[],
|
||||
0.18,
|
||||
True,
|
||||
),
|
||||
(
|
||||
[
|
||||
{"stage0_tokens": {"output_token_count": 3, "tpot_ms": 0.0}},
|
||||
{"stage0_tokens": {"output_token_count": 2, "tpot_ms": 0.0}},
|
||||
],
|
||||
[],
|
||||
0.1,
|
||||
False,
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_adapter_reports_exact_or_weighted_token_timing(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
request_metrics: list[dict[str, object]],
|
||||
expected_itl: list[float],
|
||||
expected_text_latency: float,
|
||||
expected_tpot_measured: bool,
|
||||
):
|
||||
case, options = _case(tmp_path), _session_options(tmp_path)
|
||||
request = _request(case, options)
|
||||
request.omniinteract_prepared_input = data.OmniInteractPreparedInput(
|
||||
1.0, b"pcm", ("frame",), "data:audio/wav;base64,ref"
|
||||
)
|
||||
|
||||
async def run(*args, **kwargs):
|
||||
return oi.OmniInteractCaseResult(
|
||||
case.subset,
|
||||
str(case.video_path),
|
||||
str(options.output_root),
|
||||
success=True,
|
||||
output_tokens=5,
|
||||
duplex_request_metrics=request_metrics,
|
||||
duplex_session_metrics={"mean_ttft_ms": 100.0},
|
||||
)
|
||||
|
||||
monkeypatch.setattr(benchmark_patch, "run_omniinteract_case", run)
|
||||
output = await benchmark_patch.async_request_openai_realtime_duplex(request, session=None)
|
||||
|
||||
assert output.success
|
||||
assert output.itl == expected_itl
|
||||
assert output.text_latency == pytest.approx(expected_text_latency)
|
||||
assert output.tpot_measured is expected_tpot_measured
|
||||
|
||||
|
||||
def test_batch_finalization_publishes_only_measured_results(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
case, options = _case(tmp_path), _session_options(tmp_path)
|
||||
options.output_root.mkdir()
|
||||
sample = _sample(case, options)
|
||||
result = oi.OmniInteractCaseResult(
|
||||
case.subset, str(case.video_path), str(options.output_root), success=True, eligible_for_official_eval=True
|
||||
)
|
||||
output = benchmark_patch.MixRequestFuncOutput(success=True)
|
||||
output.omniinteract_case_result = result
|
||||
published = []
|
||||
monkeypatch.setattr(benchmark_patch, "publish_deferred_case_artifacts", lambda *args: published.append(args))
|
||||
summary = benchmark_patch._finalize_omniinteract_batch([sample], [output])
|
||||
assert summary and summary["total"] == summary["success"] == summary["eligible_for_official_eval"] == 1
|
||||
assert summary["failed"] == summary["successful_but_ineligible"] == summary["audio_clipped_bytes"] == 0
|
||||
assert published == [(options.output_root, case, result)]
|
||||
|
||||
|
||||
def test_case_artifact_failures_do_not_discard_serving_metrics(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
case, options = _case(tmp_path), _session_options(tmp_path)
|
||||
sample = _sample(case, options)
|
||||
result = oi.OmniInteractCaseResult(
|
||||
case.subset, str(case.video_path), str(options.output_root), success=True, eligible_for_official_eval=True
|
||||
)
|
||||
output = benchmark_patch.MixRequestFuncOutput(success=True, error="inference warning")
|
||||
output.omniinteract_case_result = result
|
||||
monkeypatch.setattr(
|
||||
benchmark_patch,
|
||||
"publish_deferred_case_artifacts",
|
||||
lambda *args: (_ for _ in ()).throw(OSError("case disk full")),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
benchmark_patch,
|
||||
"write_failure_artifacts",
|
||||
lambda *args: (_ for _ in ()).throw(OSError("failure disk full")),
|
||||
)
|
||||
monkeypatch.setattr(benchmark_patch, "write_omniinteract_batch_artifacts", lambda *args: None)
|
||||
|
||||
summary = benchmark_patch._finalize_omniinteract_batch([sample], [output])
|
||||
metrics, _ = calculate_metrics(
|
||||
[],
|
||||
[output],
|
||||
1.0,
|
||||
None,
|
||||
[50.0],
|
||||
{},
|
||||
TaskType.GENERATION,
|
||||
[],
|
||||
None,
|
||||
float("inf"),
|
||||
1.0,
|
||||
)
|
||||
|
||||
assert summary and summary["failed"] == 1
|
||||
assert summary["artifacts_complete"] is False
|
||||
assert len(summary["artifact_errors"]) == 2
|
||||
assert output.success and metrics.completed == 1
|
||||
assert output.error.startswith("inference warning\n")
|
||||
assert "Artifact publication failed: case disk full" in output.error
|
||||
|
||||
|
||||
def test_batch_artifact_failure_removes_partial_batch_files(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
case, options = _case(tmp_path), _session_options(tmp_path)
|
||||
sample = _sample(case, options)
|
||||
result = oi.OmniInteractCaseResult(
|
||||
case.subset, str(case.video_path), str(options.output_root), success=True, eligible_for_official_eval=True
|
||||
)
|
||||
output = benchmark_patch.MixRequestFuncOutput(success=True)
|
||||
output.omniinteract_case_result = result
|
||||
monkeypatch.setattr(benchmark_patch, "publish_deferred_case_artifacts", lambda *args: None)
|
||||
|
||||
def fail_batch(root, *args):
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
(root / "batch_summary.json").write_text("partial")
|
||||
raise OSError("batch disk full")
|
||||
|
||||
monkeypatch.setattr(benchmark_patch, "write_omniinteract_batch_artifacts", fail_batch)
|
||||
|
||||
summary = benchmark_patch._finalize_omniinteract_batch([sample], [output])
|
||||
|
||||
assert summary and summary["success"] == 1 and summary["artifacts_complete"] is False
|
||||
assert summary["artifact_errors"] == ["Batch artifact publication failed: batch disk full"]
|
||||
assert output.success
|
||||
assert not (options.output_root / "batch_summary.json").exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "value", "match"),
|
||||
[
|
||||
("omniinteract_ref_audio", None, "ref-audio"),
|
||||
("endpoint", "/v1/completions", "v1/realtime"),
|
||||
("max_concurrency", 0, "positive"),
|
||||
("skip_tokenizer_init", True, "skip-tokenizer-init"),
|
||||
],
|
||||
)
|
||||
def test_cli_rejects_incompatible_omniinteract_options(field: str, value, match: str):
|
||||
args = argparse.Namespace(
|
||||
dataset_name="omniinteract",
|
||||
backend="openai-realtime-duplex",
|
||||
endpoint="/v1/realtime",
|
||||
omniinteract_ref_audio="ref.wav",
|
||||
max_concurrency=1,
|
||||
skip_tokenizer_init=False,
|
||||
)
|
||||
setattr(args, field, value)
|
||||
with pytest.raises(ValueError, match=match):
|
||||
preprocess_serve_args(args)
|
||||
|
||||
|
||||
def test_whole_benchmark_lock_serializes_a_shared_output_root(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
entered, release = threading.Event(), threading.Event()
|
||||
calls: list[int] = []
|
||||
|
||||
async def fake_main_async(args):
|
||||
calls.append(len(calls))
|
||||
if len(calls) == 1:
|
||||
entered.set()
|
||||
await asyncio.to_thread(release.wait)
|
||||
return {"total": 1}
|
||||
|
||||
monkeypatch.setattr(benchmark_serve, "main_async", fake_main_async)
|
||||
args = argparse.Namespace(dataset_name="omniinteract", omniinteract_output_dir=tmp_path, extra_body=None)
|
||||
with ThreadPoolExecutor(2) as executor:
|
||||
first = executor.submit(benchmark_serve.main, args)
|
||||
assert entered.wait(1)
|
||||
second = executor.submit(benchmark_serve.main, args)
|
||||
time.sleep(0.05)
|
||||
assert calls == [0]
|
||||
release.set()
|
||||
assert first.result() == second.result() == {"total": 1}
|
||||
assert calls == [0, 1]
|
||||
@@ -1,3 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
@@ -22,6 +25,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
|
||||
"dataset_name",
|
||||
[
|
||||
"daily-omni",
|
||||
"omniinteract",
|
||||
"seed-tts",
|
||||
"seed-tts-text",
|
||||
"seed-tts-design",
|
||||
@@ -55,6 +59,7 @@ def test_extend_omni_choices_updates_tracking_parser_shadow(dataset_name: str) -
|
||||
[
|
||||
(["--print-stage"], "print_stage", True),
|
||||
(["--daily-omni-input-mode", "audio"], "daily_omni_input_mode", "audio"),
|
||||
(["--omniinteract-subsets", "1qna"], "omniinteract_subsets", ["1qna"]),
|
||||
(["--seed-tts-locale", "zh"], "seed_tts_locale", "zh"),
|
||||
],
|
||||
)
|
||||
@@ -80,12 +85,14 @@ def test_add_omni_args_preserves_implicit_defaults() -> None:
|
||||
args = parser.parse_args([])
|
||||
assert args.print_stage is False
|
||||
assert args.daily_omni_input_mode == "all"
|
||||
assert args.omniinteract_subsets == ["1q1a", "1q1a_math", "1qna"]
|
||||
assert args.seed_tts_locale == "en"
|
||||
assert args.explicit_keys == set()
|
||||
|
||||
|
||||
def test_update_omni_help_updates_upstream_actions() -> None:
|
||||
parser = TrackingArgumentParser()
|
||||
parser.add_argument("--num-prompts", default=1000, help="Number of prompts.")
|
||||
parser.add_argument("--percentile-metrics", help="Upstream percentile help.")
|
||||
parser.add_argument("--random-mm-limit-mm-per-prompt", help="Upstream limit help.")
|
||||
parser.add_argument("--random-mm-bucket-config", help="Upstream bucket help.")
|
||||
@@ -93,6 +100,7 @@ def test_update_omni_help_updates_upstream_actions() -> None:
|
||||
update_omni_help(parser)
|
||||
|
||||
actions = {action.dest: action for action in parser._actions}
|
||||
assert "OmniInteract uses 3" in actions["num_prompts"].help
|
||||
assert all(metric in actions["percentile_metrics"].help for metric in ("ttfc", "tpop", "audio_rtf"))
|
||||
assert "probabilities are renormalized" in actions["random_mm_limit_mm_per_prompt"].help
|
||||
assert "Currently allows for 3 modalities" in actions["random_mm_bucket_config"].help
|
||||
@@ -120,6 +128,31 @@ def test_preprocess_serve_args_merges_bot_task_without_overriding_extra_body(
|
||||
assert args.extra_body == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("argv", "expected"),
|
||||
[
|
||||
([], 3),
|
||||
(["--num-prompts", "0"], 0),
|
||||
(["--num-prompts", "12"], 12),
|
||||
],
|
||||
)
|
||||
def test_preprocess_serve_args_applies_safe_omniinteract_prompt_default(
|
||||
argv: list[str],
|
||||
expected: int,
|
||||
) -> None:
|
||||
parser = TrackingArgumentParser()
|
||||
parser.add_argument("--dataset-name", default="omniinteract")
|
||||
parser.add_argument("--backend", default="openai-realtime-duplex")
|
||||
parser.add_argument("--endpoint", default="/v1/realtime")
|
||||
parser.add_argument("--omniinteract-ref-audio", default="reference.wav")
|
||||
parser.add_argument("--num-prompts", type=int, default=1000)
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
preprocess_serve_args(args)
|
||||
|
||||
assert args.num_prompts == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("argv", "expected_extra_body", "expected_explicit"),
|
||||
[
|
||||
|
||||
@@ -1,10 +1,12 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2] / ".buildkite" / "common" / "scripts"))
|
||||
|
||||
@@ -38,8 +40,6 @@ def _render(changed_files: list[str]) -> str:
|
||||
|
||||
|
||||
def test_bootstrap_if_injected_by_step_key() -> None:
|
||||
import yaml
|
||||
|
||||
rendered = _render_bootstrap_pipeline(
|
||||
BOOTSTRAP_STEPS_TEMPLATE,
|
||||
decision=resolve_ci_decision([]),
|
||||
@@ -62,16 +62,23 @@ def test_bootstrap_steps_loaded_from_file() -> None:
|
||||
|
||||
|
||||
def test_docs_only_allows_main_scheduled_nightly_weekly_only() -> None:
|
||||
"""skip_all: no PR labels; main + NIGHTLY=1 / WEEKLY=1 still gates scheduled CI."""
|
||||
"""skip_all: no PR labels; main + NIGHTLY=1 / WEEKLY=1 / NON_CRITICAL=1 still gates scheduled CI."""
|
||||
rendered = _render(["docs/foo.md"])
|
||||
assert "key: image-build" in rendered
|
||||
assert "key: upload-nightly-pipeline" in rendered
|
||||
assert "key: upload-weekly-pipeline" in rendered
|
||||
# Scheduled nightly also uploads L2/L3 with --e2e
|
||||
assert "key: upload-ready-pipeline" in rendered
|
||||
assert "key: upload-merge-pipeline" in rendered
|
||||
assert 'build.env("NIGHTLY") == "1"' in rendered
|
||||
# Scheduled main+WEEKLY=1 uploads L2/L3 with --e2e; NIGHTLY still gates L4 only.
|
||||
doc = yaml.safe_load(rendered)
|
||||
by_key = {step["key"]: step for step in doc["steps"]}
|
||||
assert "NIGHTLY" not in by_key["upload-ready-pipeline"]["if"]
|
||||
assert 'build.branch == "main"' in by_key["upload-ready-pipeline"]["if"]
|
||||
assert 'build.env("WEEKLY") == "1"' in by_key["upload-ready-pipeline"]["if"]
|
||||
assert "NIGHTLY" not in by_key["upload-merge-pipeline"]["if"]
|
||||
assert 'build.branch == "main"' in by_key["upload-merge-pipeline"]["if"]
|
||||
assert 'build.env("WEEKLY") == "1"' in by_key["upload-merge-pipeline"]["if"]
|
||||
assert 'build.env("NIGHTLY") == "1"' in by_key["upload-nightly-pipeline"]["if"]
|
||||
assert 'build.env("WEEKLY") == "1"' in rendered
|
||||
assert 'build.env("NON_CRITICAL") == "1"' in rendered
|
||||
assert "nightly-test" not in rendered
|
||||
assert "weekly-test" not in rendered
|
||||
assert "merge-test" not in rendered
|
||||
@@ -79,6 +86,19 @@ def test_docs_only_allows_main_scheduled_nightly_weekly_only() -> None:
|
||||
assert "if: false" not in rendered
|
||||
|
||||
|
||||
def test_npu_docs_only_does_not_upload_ready_on_nightly() -> None:
|
||||
"""NPU skip_all: scheduled NIGHTLY still uploads L4, not L2 ready."""
|
||||
rendered = _render_bootstrap_pipeline(
|
||||
BOOTSTRAP_STEPS_TEMPLATE,
|
||||
decision=resolve_ci_decision(["docs/foo.md"]),
|
||||
path=Path(".buildkite/npu/bootstrap-upload-steps.yml"),
|
||||
)
|
||||
doc = yaml.safe_load(rendered)
|
||||
by_key = {step["key"]: step for step in doc["steps"]}
|
||||
assert "upload-ready-pipeline" not in by_key
|
||||
assert 'build.env("NIGHTLY") == "1"' in by_key["upload-nightly-pipeline"]["if"]
|
||||
|
||||
|
||||
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
|
||||
@@ -139,3 +159,65 @@ def test_mirror_hardwares_a2b3_npu_4_expands_agents_image_and_plugins() -> None:
|
||||
assert step["plugins"][0]["kubernetes"]["podSpecPatch"]["imagePullSecrets"] == [
|
||||
{"name": "swr-secret"},
|
||||
]
|
||||
|
||||
|
||||
def _surviving_labels(doc: dict, changed_files: list[str]) -> set[str]:
|
||||
rendered = _render_test_pipeline(doc, changed_files=changed_files)
|
||||
labels: set[str] = set()
|
||||
|
||||
def walk(steps: list | None) -> None:
|
||||
for step in steps or []:
|
||||
if not isinstance(step, dict):
|
||||
continue
|
||||
if "label" in step:
|
||||
labels.add(step["label"])
|
||||
walk(step.get("steps"))
|
||||
|
||||
walk(rendered.get("steps"))
|
||||
return labels
|
||||
|
||||
|
||||
# Synthetic coverage-style job: shared inputs that change what the split measures.
|
||||
_COVERAGE_SHARED_INPUTS_DOC = {
|
||||
"steps": [
|
||||
{
|
||||
"label": "Coverage Pilot",
|
||||
"source_file_dependencies": [
|
||||
"tests/e2e/online_serving/test_example.py",
|
||||
".buildkite/common/scripts/run_cov_split.sh",
|
||||
"pyproject.toml",
|
||||
],
|
||||
"commands": [".buildkite/common/scripts/run_cov_split.sh --model-id example"],
|
||||
},
|
||||
{
|
||||
"label": "Unrelated Model Test",
|
||||
"source_file_dependencies": [
|
||||
"tests/e2e/online_serving/test_other.py",
|
||||
],
|
||||
"commands": ["pytest -sv tests/e2e/online_serving/test_other.py"],
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"changed_file",
|
||||
[
|
||||
".buildkite/common/scripts/run_cov_split.sh",
|
||||
"pyproject.toml",
|
||||
],
|
||||
)
|
||||
def test_coverage_shared_inputs_select_dependent_job(changed_file: str) -> None:
|
||||
"""Jobs that list coverage shared inputs must stay selected when those files change."""
|
||||
labels = _surviving_labels(_COVERAGE_SHARED_INPUTS_DOC, [changed_file])
|
||||
assert "Coverage Pilot" in labels
|
||||
assert "Unrelated Model Test" not in labels
|
||||
|
||||
|
||||
def test_coverage_shared_inputs_ignored_for_unrelated_change() -> None:
|
||||
labels = _surviving_labels(
|
||||
_COVERAGE_SHARED_INPUTS_DOC,
|
||||
["vllm_omni/entrypoints/openai/serving_chat.py"],
|
||||
)
|
||||
assert "Coverage Pilot" not in labels
|
||||
assert "Unrelated Model Test" not in labels
|
||||
@@ -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 StageConfigFactory and related classes.
|
||||
"""
|
||||
@@ -14,7 +14,8 @@ import pytest
|
||||
from transformers import PretrainedConfig, Qwen3OmniMoeConfig
|
||||
|
||||
from tests.helpers.stage_config import get_deploy_config_path, get_deploy_config_stage
|
||||
from vllm_omni.config.config_factory import StageConfigFactory
|
||||
from vllm_omni.config import config_factory as config_factory_module
|
||||
from vllm_omni.config.config_factory import StageConfigFactory, _materialize_object_storage_configs
|
||||
from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability
|
||||
from vllm_omni.config.omni_config import VllmOmniConfig
|
||||
from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline, resolve_pipeline_config
|
||||
@@ -33,6 +34,7 @@ from vllm_omni.config.stage_config import (
|
||||
build_stage_runtime_overrides,
|
||||
load_deploy_config,
|
||||
merge_pipeline_deploy,
|
||||
normalize_pipeline_cli_overrides,
|
||||
pipeline_cfg_resolver,
|
||||
)
|
||||
from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
|
||||
@@ -55,6 +57,7 @@ def clear_config_factory_caches():
|
||||
yield
|
||||
StageConfigFactory.get_hf_config.cache_clear()
|
||||
StageConfigFactory.try_infer_model_type.cache_clear()
|
||||
_materialize_object_storage_configs.cache_clear()
|
||||
|
||||
|
||||
Q3_OMNI_ALL_STAGES_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=True)
|
||||
@@ -427,6 +430,20 @@ class TestStageConfigFactory:
|
||||
assert cfg["final_output"] is True
|
||||
assert cfg["final_output_type"] == "image"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_class_name",
|
||||
[
|
||||
"MiniMaxH3Pipeline",
|
||||
"WanPipeline",
|
||||
"LTX2Pipeline",
|
||||
"Cosmos3OmniDiffusersPipeline",
|
||||
],
|
||||
)
|
||||
def test_default_video_diffusion_output_type(self, model_class_name):
|
||||
configs = StageConfigFactory.create_default_diffusion({"model_class_name": model_class_name})
|
||||
|
||||
assert configs[0]["final_output_type"] == "video"
|
||||
|
||||
def test_default_diffusion_with_parallel_config(self):
|
||||
"""Test diffusion config calculates devices from parallel_config."""
|
||||
|
||||
@@ -532,6 +549,7 @@ class TestPipelineDiscovery:
|
||||
"""Check that specific models are in OMNI_PIPELINES."""
|
||||
assert "qwen2_5_omni" in OMNI_PIPELINES
|
||||
assert "qwen3_omni_moe" in OMNI_PIPELINES
|
||||
assert "qwen3_omni_moe_thinker_only" in OMNI_PIPELINES
|
||||
assert "qwen3_tts" in OMNI_PIPELINES
|
||||
|
||||
def test_registry_resolver_qwen3_omni_all_stages(self):
|
||||
@@ -556,6 +574,18 @@ class TestPipelineDiscovery:
|
||||
assert pipeline.default_deploy_config_name is None
|
||||
assert len(pipeline.stages) == 1 # thinker only
|
||||
|
||||
def test_registry_qwen3_omni_thinker_only_key_is_static(self):
|
||||
"""Explicit registry key must not re-enter the HF-config resolver."""
|
||||
pipeline = resolve_pipeline_config(
|
||||
"qwen3_omni_moe_thinker_only",
|
||||
Q3_OMNI_ALL_STAGES_HF_CONFIG,
|
||||
)
|
||||
assert pipeline is OMNI_PIPELINES["qwen3_omni_moe_thinker_only"]
|
||||
assert pipeline.model_type == "qwen3_omni_moe_thinker_only"
|
||||
assert len(pipeline.stages) == 1
|
||||
assert pipeline.stages[0].engine_output_type == "text"
|
||||
assert pipeline.stages[0].final_output_type == "text"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"pipeline",
|
||||
[registered for registered in OMNI_PIPELINES.values() if isinstance(registered, PipelineConfig)],
|
||||
@@ -1100,8 +1130,42 @@ class TestResolveScheduler:
|
||||
|
||||
|
||||
class TestDeployConfigLoading:
|
||||
def test_load_minicpmo_duplex_deploy_config(self):
|
||||
deploy_path = Path(get_deploy_config_path("minicpmo_4_5_duplex.yaml"))
|
||||
def test_rejects_legacy_stage_args_schema(self, tmp_path):
|
||||
deploy_path = tmp_path / "legacy.yaml"
|
||||
deploy_path.write_text("stage_args:\n - stage_id: 0\n", encoding="utf-8")
|
||||
|
||||
with pytest.raises(ValueError, match=r"stage_args.*PipelineConfig.*stages"):
|
||||
load_deploy_config(deploy_path)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("filename", "max_sessions"),
|
||||
[
|
||||
("minicpmo_4_5.yaml", 4),
|
||||
("minicpmo_4_5_2gpu.yaml", 4),
|
||||
("minicpmo_4_5_3gpu.yaml", 4),
|
||||
("minicpmo_4_5_8x4090.yaml", 1),
|
||||
("minicpmo_4_5_3gpu_stage1_replicas.yaml", 4),
|
||||
("minicpmo_4_5_4gpu_stage1_replicas.yaml", 4),
|
||||
("minicpmo_4_5_8x4090_stage1_replicas.yaml", 1),
|
||||
],
|
||||
)
|
||||
def test_minicpmo_deploy_configs_enable_duplex(self, filename: str, max_sessions: int):
|
||||
"""Every MiniCPM-o profile serves /v1/realtime, so the retired
|
||||
minicpmo_4_5_duplex.yaml overlay has no replacement to fall back on.
|
||||
"""
|
||||
deploy = load_deploy_config(Path(get_deploy_config_path(filename)))
|
||||
pipeline = resolve_pipeline_config("minicpmo_4_5")
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
assert deploy.session_mode == "duplex"
|
||||
assert deploy.active_stream_window == max_sessions
|
||||
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"]
|
||||
|
||||
def test_load_minicpmo_default_deploy_config(self):
|
||||
deploy_path = Path(get_deploy_config_path("minicpmo_4_5.yaml"))
|
||||
|
||||
deploy = load_deploy_config(deploy_path)
|
||||
pipeline = resolve_pipeline_config("minicpmo_4_5")
|
||||
@@ -1109,18 +1173,18 @@ class TestDeployConfigLoading:
|
||||
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
assert deploy.session_mode == "duplex"
|
||||
assert deploy.async_chunk is True
|
||||
assert deploy.active_stream_window == 1
|
||||
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 [stage.yaml_engine_args["async_scheduling"] for stage in stages] == [False, False, False]
|
||||
assert all("Async" not in (stage.scheduler_cls or "") for stage in stages)
|
||||
assert stages[0].yaml_engine_args["async_scheduling"] is True
|
||||
assert stages[1].yaml_engine_args["async_scheduling"] is True
|
||||
assert "Async" in (stages[0].scheduler_cls or "")
|
||||
assert "Async" in (stages[1].scheduler_cls or "")
|
||||
assert [stage.devices for stage in deploy.stages] == ["0", "0", "0"]
|
||||
assert deploy.stages[1].enforce_eager is False
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 2048
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["min_tokens"] == 0
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [1]
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 4096
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["min_tokens"] == 50
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["temperature"] == 0.8
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [6561]
|
||||
assert "codec_sampling_params" not in stages[1].yaml_engine_args
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("filename", "stage0_devices", "stage1_devices", "stage2_devices", "stage1_replicas"),
|
||||
@@ -1189,14 +1253,22 @@ stages:
|
||||
assert deploy.connectors is not None
|
||||
assert deploy.platforms is not None
|
||||
|
||||
def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self):
|
||||
@pytest.mark.parametrize(
|
||||
("hf_config", "model"),
|
||||
[
|
||||
(Q3_OMNI_THINKER_HF_CONFIG, "Qwen/Qwen3-Omni-30B-A3B-Thinking"),
|
||||
(Q3_OMNI_ALL_STAGES_HF_CONFIG, "Qwen/Qwen3-Omni-30B-A3B-Instruct"),
|
||||
],
|
||||
ids=["thinker_checkpoint", "instruct_forced_thinker_only"],
|
||||
)
|
||||
def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self, hf_config, model):
|
||||
deploy_path = Path(get_deploy_config_path("qwen3_omni_moe_thinking.yaml"))
|
||||
with patch(
|
||||
"vllm_omni.config.config_factory.get_config",
|
||||
return_value=Q3_OMNI_THINKER_HF_CONFIG,
|
||||
return_value=hf_config,
|
||||
):
|
||||
pipeline = StageConfigFactory.get_pipeline_config(
|
||||
"Qwen/Qwen3-Omni-30B-A3B-Thinking",
|
||||
model,
|
||||
trust_remote_code=False,
|
||||
deploy_config_path=str(deploy_path),
|
||||
)
|
||||
@@ -1205,6 +1277,7 @@ stages:
|
||||
assert pipeline.model_type == "qwen3_omni_moe_thinker_only"
|
||||
|
||||
deploy = load_deploy_config(deploy_path)
|
||||
assert deploy.pipeline == "qwen3_omni_moe_thinker_only"
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
assert len(stages) == 1
|
||||
@@ -1266,6 +1339,13 @@ stages:
|
||||
("step_audio2_ci.yaml", "step_audio_2_asr", 1, "text", "step_audio_2_asr"),
|
||||
("hunyuan_video_15.yaml", "hunyuan_video_15", 1, "video", None),
|
||||
("wan2_2_ti2v.yaml", "wan2_2_ti2v", 1, "video", None),
|
||||
(
|
||||
"minimax_h3_disaggregated.yaml",
|
||||
"minimax_h3_disaggregated",
|
||||
2,
|
||||
"video",
|
||||
"minimax_h3_disaggregated",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_load_new_registry_backed_deploy_configs(
|
||||
@@ -1286,15 +1366,98 @@ stages:
|
||||
assert len(stages) == stage_count
|
||||
assert stages[-1].final_output is True
|
||||
assert stages[-1].final_output_type == final_output_type
|
||||
if pipeline_name == "minimax_h3_disaggregated":
|
||||
assert stages[0].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_model_root")
|
||||
assert (
|
||||
stages[1].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_diffusion_model_path")
|
||||
)
|
||||
assert "model_path_resolver" not in stages[0].yaml_extras
|
||||
assert "model_path_resolver" not in stages[1].yaml_extras
|
||||
if deploy.trust_remote_code is not None:
|
||||
assert {s.yaml_engine_args.get("trust_remote_code") for s in stages} == {deploy.trust_remote_code}
|
||||
|
||||
def test_minimax_h3_resolvers_reach_production_engine_args(self, tmp_path: Path, monkeypatch):
|
||||
from vllm_omni.engine import stage_init_utils
|
||||
|
||||
model_root = tmp_path / "MiniMax-H3"
|
||||
(model_root / "FL2VA" / "text_encoder").mkdir(parents=True)
|
||||
(model_root / "Ref2VA" / "text_encoder").mkdir(parents=True)
|
||||
monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda _engine_args: None)
|
||||
|
||||
deploy = load_deploy_config(get_deploy_config_path("minimax_h3_disaggregated.yaml"))
|
||||
assert deploy.stages[1].engine_extras["model_loaded"] == {"text_encoder": False}
|
||||
stages = merge_pipeline_deploy(OMNI_PIPELINES["minimax_h3_disaggregated"], deploy)
|
||||
resolved = [stage_init_utils.build_engine_args_dict(stage.to_omegaconf(), str(model_root)) for stage in stages]
|
||||
|
||||
assert resolved[0]["model"] == str(model_root / "FL2VA" / "text_encoder")
|
||||
assert resolved[1]["model"] == str(model_root / "FL2VA")
|
||||
assert all("model_path_resolver" not in engine_args for engine_args in resolved)
|
||||
|
||||
def test_minimax_h3_text_encoder_tp_reaches_only_stage_zero(self):
|
||||
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
|
||||
stages, _ = StageConfigFactory._create_legacy_from_registry(
|
||||
pipeline,
|
||||
{"text_encoder_tp_size": 4},
|
||||
)
|
||||
|
||||
assert stages[0].runtime_overrides["tensor_parallel_size"] == 4
|
||||
assert "text_encoder_tp_size" not in stages[0].runtime_overrides
|
||||
assert "text_encoder_tp_size" not in stages[1].runtime_overrides
|
||||
assert stages[1].yaml_engine_args["parallel_config"]["tensor_parallel_size"] == 1
|
||||
|
||||
def test_minimax_h3_disaggregated_defaults_match_validated_topology(self):
|
||||
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
|
||||
deploy = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated.yaml")))
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
assert stages[0].yaml_engine_args["max_num_seqs"] == 1
|
||||
assert stages[0].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_model_root")
|
||||
assert stages[1].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_diffusion_model_path")
|
||||
parallel = stages[1].yaml_engine_args["parallel_config"]
|
||||
assert parallel["tensor_parallel_size"] == 1
|
||||
assert parallel["ulysses_degree"] == 4
|
||||
assert parallel["vae_patch_parallel_size"] == 4
|
||||
assert stages[1].yaml_engine_args["inline_diffusion"] is True
|
||||
assert stages[1].yaml_extras["default_sampling_params"]["num_inference_steps"] == 50
|
||||
|
||||
turbo = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated_turbo.yaml")))
|
||||
turbo_stages = merge_pipeline_deploy(pipeline, turbo)
|
||||
turbo_sampling = turbo_stages[1].yaml_extras["default_sampling_params"]
|
||||
assert turbo_sampling["num_inference_steps"] == 5
|
||||
assert turbo_sampling["extra_args"] == {"flow_shift": 6.0, "audio_flow_shift": 3.0}
|
||||
|
||||
def test_minimax_h3_text_encoder_tp_alias_targets_stage_zero(self):
|
||||
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
|
||||
|
||||
assert normalize_pipeline_cli_overrides(pipeline, {}) == {}
|
||||
assert normalize_pipeline_cli_overrides(pipeline, {"text_encoder_tp_size": 4}) == {
|
||||
"stage_0_tensor_parallel_size": 4
|
||||
}
|
||||
|
||||
def test_minimax_h3_stage_zero_tp_override_wins_alias_conflict(self):
|
||||
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
|
||||
|
||||
with pytest.warns(UserWarning, match="stage_0_tensor_parallel_size=2 takes precedence"):
|
||||
normalized = normalize_pipeline_cli_overrides(
|
||||
pipeline,
|
||||
{"text_encoder_tp_size": 4, "stage_0_tensor_parallel_size": 2},
|
||||
)
|
||||
|
||||
assert normalized == {"stage_0_tensor_parallel_size": 2}
|
||||
|
||||
def test_minimax_h3_rejects_text_encoder_tp_on_diffusion_stage(self):
|
||||
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
|
||||
|
||||
with pytest.raises(ValueError, match="stage_1_text_encoder_tp_size cannot be set"):
|
||||
normalize_pipeline_cli_overrides(pipeline, {"stage_1_text_encoder_tp_size": 4})
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("config_json", "model_index", "expected_pipeline"),
|
||||
[
|
||||
({"model_type": "step_audio_2"}, None, "step_audio_2"),
|
||||
(None, {"_class_name": "HunyuanVideo15Pipeline"}, "hunyuan_video_15"),
|
||||
(None, {"_class_name": "WanPipeline"}, "wan2_2_ti2v"),
|
||||
(None, {"_class_name": "WanDMDPipeline"}, "wan2_2_ti2v"),
|
||||
],
|
||||
)
|
||||
def test_migrated_models_are_discovered_without_explicit_deploy(
|
||||
@@ -1323,6 +1486,29 @@ stages:
|
||||
assert pipeline is not None
|
||||
assert pipeline.model_type == expected_pipeline
|
||||
|
||||
def test_minimax_h3_disaggregation_is_not_auto_discovered(self):
|
||||
def get_model_file(filename, _model, revision=None):
|
||||
del revision
|
||||
if filename == "config.json":
|
||||
return None
|
||||
if filename == "model_index.json":
|
||||
return {"_class_name": "MiniMaxH3Pipeline"}
|
||||
return None
|
||||
|
||||
with (
|
||||
patch.object(StageConfigFactory, "get_hf_config", return_value=None),
|
||||
patch(
|
||||
"vllm_omni.config.config_factory.get_hf_file_to_dict",
|
||||
side_effect=get_model_file,
|
||||
),
|
||||
):
|
||||
pipeline = StageConfigFactory.get_pipeline_config(
|
||||
model="/models/unrelated-checkpoint-name",
|
||||
trust_remote_code=False,
|
||||
)
|
||||
|
||||
assert pipeline is None
|
||||
|
||||
@pytest.mark.parametrize("deploy_name", ["step_audio_2.yaml", "step_audio_2_async_chunk.yaml"])
|
||||
def test_step_audio2_deploy_configs_fit_two_gpus(self, deploy_name):
|
||||
deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name)))
|
||||
@@ -1375,7 +1561,8 @@ stages:
|
||||
assert async_stages[1].custom_process_input_func is None
|
||||
|
||||
def test_no_bundled_legacy_stage_config_yamls(self):
|
||||
stage_config_dir = Path(__file__).parent.parent / "vllm_omni" / "model_executor" / "stage_configs"
|
||||
repo_root = Path(__file__).resolve().parents[2]
|
||||
stage_config_dir = repo_root / "vllm_omni" / "model_executor" / "stage_configs"
|
||||
assert not list(stage_config_dir.glob("*.yaml"))
|
||||
|
||||
def test_merge_pipeline_deploy(self):
|
||||
@@ -1649,6 +1836,7 @@ stages:
|
||||
]
|
||||
}
|
||||
original_foo = fake_config["stages"][0]["foo"]
|
||||
assert isinstance(original_foo, dict)
|
||||
original_b = original_foo["b"]
|
||||
|
||||
with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config):
|
||||
@@ -2101,6 +2289,26 @@ class TestMingFlashOmniPipeline:
|
||||
class TestBaseConfigInheritance:
|
||||
"""Test deploy YAML base_config inheritance."""
|
||||
|
||||
def test_minicpmo_overlays_inherit_talker_sampling_params(self):
|
||||
"""Overlays must keep the base Talker codec Sampler knobs."""
|
||||
for filename in (
|
||||
"minicpmo_4_5_2gpu.yaml",
|
||||
"minicpmo_4_5_3gpu.yaml",
|
||||
"minicpmo_4_5_3gpu_stage1_replicas.yaml",
|
||||
"minicpmo_4_5_4gpu_stage1_replicas.yaml",
|
||||
"minicpmo_4_5_8x4090_stage1_replicas.yaml",
|
||||
):
|
||||
path = Path(get_deploy_config_path(filename))
|
||||
if not path.exists():
|
||||
pytest.skip(f"{filename} not found")
|
||||
deploy = load_deploy_config(path)
|
||||
sampling = deploy.stages[1].default_sampling_params
|
||||
assert sampling is not None, f"{filename} stage 1 lost default_sampling_params"
|
||||
assert sampling["temperature"] == 0.8, filename
|
||||
assert sampling["top_k"] == 25, filename
|
||||
assert sampling["top_p"] == 0.85, filename
|
||||
assert sampling["repetition_penalty"] == 1.05, filename
|
||||
|
||||
def test_ci_inherits_from_main(self):
|
||||
ci_path = Path(get_deploy_config_path("ci/qwen3_omni_moe.yaml"))
|
||||
if not ci_path.exists():
|
||||
@@ -2191,6 +2399,42 @@ class TestPlatformOverrides:
|
||||
assert rocm.stages[0].enforce_eager is None
|
||||
assert rocm.stages[1].enforce_eager is True
|
||||
|
||||
def test_higgs_audio_v3_rocm_uses_triton_attention(self):
|
||||
deploy_path = Path(get_deploy_config_path("higgs_multimodal_qwen3.yaml"))
|
||||
|
||||
base = load_deploy_config(deploy_path)
|
||||
assert base.stages[0].engine_extras["attention_backend"] == "FLASHINFER"
|
||||
|
||||
rocm = _apply_platform_overrides(base, platform="rocm")
|
||||
assert rocm.stages[0].engine_extras["attention_backend"] == "TRITON_ATTN"
|
||||
|
||||
pipeline = resolve_pipeline_config("higgs_multimodal_qwen3")
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
stages = merge_pipeline_deploy(pipeline, rocm)
|
||||
assert stages[0].yaml_engine_args["attention_backend"] == "TRITON_ATTN"
|
||||
|
||||
def test_qwen3_omni_cuda_uses_thinker_rotary_custom_op(self):
|
||||
deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
|
||||
pipeline = resolve_pipeline_config(
|
||||
"qwen3_omni_moe",
|
||||
Q3_OMNI_ALL_STAGES_HF_CONFIG,
|
||||
)
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
|
||||
cuda = _apply_platform_overrides(load_deploy_config(deploy_path), platform="cuda")
|
||||
cuda_stages = merge_pipeline_deploy(pipeline, cuda)
|
||||
expected = {"custom_ops": ["+rotary_embedding"]}
|
||||
assert cuda_stages[0].yaml_engine_args["compilation_config"] == expected
|
||||
assert all("compilation_config" not in stage.yaml_engine_args for stage in cuda_stages[1:])
|
||||
|
||||
for platform in ("cpu", "musa", "npu", "rocm", "xpu"):
|
||||
deploy = _apply_platform_overrides(
|
||||
load_deploy_config(deploy_path),
|
||||
platform=platform,
|
||||
)
|
||||
config = deploy.stages[0].compilation_config or {}
|
||||
assert "+rotary_embedding" not in config.get("custom_ops", [])
|
||||
|
||||
def test_minicpmo_4_5_cuda_caps_talker_kv_cache(self):
|
||||
pipeline = resolve_pipeline_config("minicpmo_4_5")
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
@@ -2697,22 +2941,23 @@ class TestSentinelDefaultPrecedence:
|
||||
)
|
||||
|
||||
def test_ming_flash_omni_topology(self):
|
||||
"""Guard ming_flash_omni's SIP cleanup: stage 0 has no full-payload
|
||||
producer hook (arch is not in ``_FULL_PAYLOAD_INPUT_STAGES``), and
|
||||
stage 1 uses only ``thinker2talker_token_only`` (sync_process_input_func).
|
||||
"""Guard ming_flash_omni's SIP cleanup: stage 0 has no
|
||||
``custom_process_next_stage_input_func``, and stage 1 does not set
|
||||
``requires_full_payload_input=True``. Stage 1 uses only
|
||||
``thinker2talker_token_only`` (sync_process_input_func).
|
||||
The dead ``thinker2talker`` (custom_process_input_func) was removed
|
||||
because ming_flash_omni has no async_chunk support and both functions
|
||||
called ``_build_talker_inputs`` identically.
|
||||
Merge under either async_chunk mode must not re-introduce a
|
||||
stage-0 full-payload hook."""
|
||||
Merge under either async_chunk mode must not re-introduce either
|
||||
full-payload transport gate."""
|
||||
pipeline = resolve_pipeline_config("ming_flash_omni")
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
|
||||
stage0, stage1 = pipeline.stages
|
||||
assert stage0.custom_process_next_stage_input_func is None, (
|
||||
"ming_flash_omni stage 0 must not declare a full-payload producer "
|
||||
"(connector path is not active for this arch)."
|
||||
"ming_flash_omni stage 0 must not declare a full-payload producer hook."
|
||||
)
|
||||
assert stage1.requires_full_payload_input is False
|
||||
assert stage1.custom_process_input_func is None
|
||||
assert stage1.sync_process_input_func is not None
|
||||
assert stage1.sync_process_input_func.endswith("thinker2talker_token_only")
|
||||
@@ -2765,3 +3010,90 @@ class TestPipelineConfigResolvers:
|
||||
pass
|
||||
|
||||
assert resolver(NotTheRightHfConfig()) is None
|
||||
|
||||
|
||||
class TestObjectStorageConfigResolution:
|
||||
"""Object-storage URIs must be materialized locally before any HF-style reads.
|
||||
|
||||
Regression coverage for review on the Run:AI PR: ``get_config`` and the
|
||||
config.json/model_index.json fallbacks crashed with ``HFValidationError``
|
||||
for ``s3://`` URIs, and the name-match fallback scanned the whole URI, so
|
||||
a bucket named after another pipeline could hijack pipeline selection.
|
||||
"""
|
||||
|
||||
@pytest.fixture
|
||||
def fake_object_storage(self, monkeypatch, tmp_path):
|
||||
"""Replace ObjectStorageModel with a fake that "pulls" into a tmp dir.
|
||||
|
||||
Returns a recorder: calls to the returned callable write files into
|
||||
the materialized directory, and ``recorder.pulls`` records every
|
||||
``pull_files`` invocation.
|
||||
"""
|
||||
pulls: list[tuple[str, list[str] | None, list[str] | None]] = []
|
||||
materialized_files: list[tuple[str, str]] = []
|
||||
|
||||
class FakeObjectStorageModel:
|
||||
def __init__(self, url: str):
|
||||
self.dir = str(tmp_path)
|
||||
|
||||
def pull_files(self, model_path, allow_pattern=None, ignore_pattern=None):
|
||||
pulls.append((model_path, allow_pattern, ignore_pattern))
|
||||
for name, content in materialized_files:
|
||||
(tmp_path / name).write_text(content)
|
||||
|
||||
monkeypatch.setattr(config_factory_module, "ObjectStorageModel", FakeObjectStorageModel)
|
||||
|
||||
class Recorder:
|
||||
def __init__(self) -> None:
|
||||
self.pulls = pulls
|
||||
|
||||
def set_files(self, files: list[tuple[str, str]]) -> None:
|
||||
materialized_files.clear()
|
||||
materialized_files.extend(files)
|
||||
|
||||
return Recorder()
|
||||
|
||||
def test_passthrough_for_non_uri(self, fake_object_storage):
|
||||
assert _materialize_object_storage_configs("org/model") == "org/model"
|
||||
assert _materialize_object_storage_configs("/local/model") == "/local/model"
|
||||
assert fake_object_storage.pulls == []
|
||||
|
||||
def test_materialize_pulls_configs_once_per_uri(self, fake_object_storage):
|
||||
uri = "s3://bucket/model"
|
||||
|
||||
first = _materialize_object_storage_configs(uri)
|
||||
second = _materialize_object_storage_configs(uri)
|
||||
|
||||
assert first == second
|
||||
assert len(fake_object_storage.pulls) == 1
|
||||
pulled_uri, allow_pattern, ignore_pattern = fake_object_storage.pulls[0]
|
||||
assert pulled_uri == uri
|
||||
assert "*.json" in allow_pattern
|
||||
assert ignore_pattern is None
|
||||
|
||||
def test_try_infer_model_type_reads_materialized_config(self, fake_object_storage):
|
||||
"""A misleading bucket name must not beat the real config.json content."""
|
||||
fake_object_storage.set_files(
|
||||
[
|
||||
(
|
||||
"config.json",
|
||||
'{"model_type": "qwen3_omni_moe", "architectures": ["Qwen3OmniMoeForConditionalGeneration"]}',
|
||||
)
|
||||
]
|
||||
)
|
||||
uri = "s3://qwen3-tts-models/Qwen3-Omni-30B-A3B-Instruct"
|
||||
|
||||
model = StageConfigFactory.try_infer_model_type(model=uri, trust_remote_code=False)
|
||||
|
||||
assert model == "qwen3_omni_moe"
|
||||
assert fake_object_storage.pulls # materialization happened
|
||||
|
||||
def test_name_match_fallback_scans_basename_only(self, fake_object_storage):
|
||||
# Empty config.json (CosyVoice3 style) forces the name-match fallback.
|
||||
fake_object_storage.set_files([("config.json", "{}")])
|
||||
|
||||
deceptive = "s3://qwen3-tts-models/plain-checkpoint"
|
||||
assert StageConfigFactory.try_infer_model_type(model=deceptive, trust_remote_code=False) is None
|
||||
|
||||
matching = "s3://any-bucket/my-cosyvoice3-model"
|
||||
assert StageConfigFactory.try_infer_model_type(model=matching, trust_remote_code=False) == "cosyvoice3"
|
||||
@@ -0,0 +1,121 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Guard the NPU pytest circular import through DiffusionOutput."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
|
||||
|
||||
def _top_level_import_modules(path: Path) -> set[str]:
|
||||
tree = ast.parse(path.read_text(encoding="utf-8"))
|
||||
modules: set[str] = set()
|
||||
for node in tree.body:
|
||||
if isinstance(node, ast.ImportFrom) and node.module:
|
||||
modules.add(node.module)
|
||||
elif isinstance(node, ast.Import):
|
||||
modules.update(alias.name for alias in node.names)
|
||||
return modules
|
||||
|
||||
|
||||
def _assert_isolated_import_succeeds(script: str) -> None:
|
||||
env = os.environ.copy()
|
||||
env.pop("CUDA_VISIBLE_DEVICES", None)
|
||||
env.pop("HIP_VISIBLE_DEVICES", None)
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", script],
|
||||
cwd=REPO_ROOT,
|
||||
env=env,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=120,
|
||||
check=False,
|
||||
)
|
||||
assert result.returncode == 0, result.stderr
|
||||
|
||||
|
||||
def test_config_package_does_not_eagerly_import_pipeline_registry() -> None:
|
||||
imported = _top_level_import_modules(REPO_ROOT / "vllm_omni" / "config" / "__init__.py")
|
||||
assert "vllm_omni.config.config_factory" not in imported
|
||||
assert "vllm_omni.config.pipeline_registry" not in imported
|
||||
|
||||
|
||||
def test_config_package_lazy_exports_still_resolve() -> None:
|
||||
from vllm_omni.config import StageConfigFactory, register_pipeline
|
||||
from vllm_omni.config.config_factory import StageConfigFactory as DirectStageConfigFactory
|
||||
from vllm_omni.config.pipeline_registry import register_pipeline as direct_register_pipeline
|
||||
|
||||
assert StageConfigFactory is DirectStageConfigFactory
|
||||
assert register_pipeline is direct_register_pipeline
|
||||
|
||||
|
||||
def test_npu_platform_defers_minimax_h3_encoder_patch() -> None:
|
||||
src = (REPO_ROOT / "vllm_omni" / "platforms" / "npu" / "platform.py").read_text(encoding="utf-8")
|
||||
init_start = src.index("def __init__(self) -> None:")
|
||||
runtime_start = src.index("def init_diffusion_model_runner_runtime")
|
||||
init_src = src[init_start:runtime_start]
|
||||
runtime_src = src[runtime_start:]
|
||||
assert "apply_minimax_h3_qwen3vl_patch" not in init_src
|
||||
assert "apply_minimax_h3_qwen3vl_patch" in runtime_src
|
||||
assert "apply_minimax_h3_qwen3vl_sdpa_patch" not in init_src
|
||||
assert "apply_minimax_h3_qwen3vl_sdpa_patch" in runtime_src
|
||||
assert "apply_minimax_h3_qwen3vl_swiglu_patch" in runtime_src
|
||||
|
||||
|
||||
def test_lora_config_import_does_not_load_pipeline_registry() -> None:
|
||||
"""Import LoRAConfig and diffusion.data in a fresh interpreter.
|
||||
|
||||
The NPU failure was ``data.py`` importing LoRAConfig while still loading,
|
||||
then ``vllm_omni.config`` re-entering ``pipeline_registry`` → PI0 →
|
||||
DiffusionOutput. This child process must not preload those exports.
|
||||
"""
|
||||
_assert_isolated_import_succeeds(
|
||||
"""
|
||||
import sys
|
||||
|
||||
from vllm_omni.config.lora import LoRAConfig
|
||||
|
||||
if "vllm_omni.config.pipeline_registry" in sys.modules:
|
||||
raise SystemExit("LoRAConfig import loaded pipeline_registry")
|
||||
if "vllm_omni.diffusion.models.pi0.pipeline_pi0" in sys.modules:
|
||||
raise SystemExit("LoRAConfig import loaded PI0_PIPELINE")
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput
|
||||
|
||||
if DiffusionOutput is None or LoRAConfig is None:
|
||||
raise SystemExit("expected LoRAConfig and DiffusionOutput")
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def test_pipeline_registry_does_not_import_pi0_runtime() -> None:
|
||||
"""Resolving the pi0 topology must not import its runtime pipeline.
|
||||
|
||||
``pipeline_pi0`` imports ``diffusion.data``. Importing it while
|
||||
``diffusion.data`` is only partially initialized closes the Ascend plugin's
|
||||
circular-import loop, so the registry must depend only on a lightweight
|
||||
pipeline-config module.
|
||||
"""
|
||||
_assert_isolated_import_succeeds(
|
||||
"""
|
||||
import sys
|
||||
|
||||
from vllm_omni.config.pipeline_registry import resolve_pipeline_config
|
||||
|
||||
if "vllm_omni.diffusion.models.pi0.pipeline_pi0" in sys.modules:
|
||||
raise SystemExit("pipeline_registry imported the pi0 runtime pipeline")
|
||||
|
||||
pipeline = resolve_pipeline_config("pi0")
|
||||
if pipeline is None or pipeline.model_type != "pi0":
|
||||
raise SystemExit("expected the registered pi0 pipeline config")
|
||||
"""
|
||||
)
|
||||
@@ -0,0 +1,317 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Drift checks for the reviewed environment-variable inventory."""
|
||||
|
||||
import ast
|
||||
import builtins
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.config.environment_variable_inventory import (
|
||||
ENVIRONMENT_VARIABLE_INVENTORY,
|
||||
EnvironmentVariableCategory,
|
||||
ModelEnvironmentVariableDisposition,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
_REPO_ROOT = Path(__file__).parents[2]
|
||||
_PACKAGE_ROOT = _REPO_ROOT / "vllm_omni"
|
||||
_INVENTORY_MODULE = _PACKAGE_ROOT / "config" / "environment_variable_inventory.py"
|
||||
_REFERENCE_PAGE = _REPO_ROOT / "docs" / "configuration" / "environment_variables.md"
|
||||
|
||||
|
||||
def _module_string_constants(tree: ast.Module) -> dict[str, str]:
|
||||
constants: dict[str, str] = {}
|
||||
for statement in tree.body:
|
||||
targets: list[ast.expr] = []
|
||||
value: ast.expr | None = None
|
||||
if isinstance(statement, ast.Assign):
|
||||
targets = statement.targets
|
||||
value = statement.value
|
||||
elif isinstance(statement, ast.AnnAssign):
|
||||
targets = [statement.target]
|
||||
value = statement.value
|
||||
if isinstance(value, ast.Constant) and isinstance(value.value, str):
|
||||
for target in targets:
|
||||
if isinstance(target, ast.Name):
|
||||
constants[target.id] = value.value
|
||||
return constants
|
||||
|
||||
|
||||
def _os_import_aliases(tree: ast.Module) -> tuple[set[str], set[str], set[str]]:
|
||||
os_names: set[str] = set()
|
||||
getenv_names: set[str] = set()
|
||||
environ_names: set[str] = set()
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.Import):
|
||||
for imported in node.names:
|
||||
if imported.name == "os":
|
||||
os_names.add(imported.asname or "os")
|
||||
elif isinstance(node, ast.ImportFrom) and node.module == "os":
|
||||
for imported in node.names:
|
||||
if imported.name == "getenv":
|
||||
getenv_names.add(imported.asname or imported.name)
|
||||
elif imported.name == "environ":
|
||||
environ_names.add(imported.asname or imported.name)
|
||||
return os_names, getenv_names, environ_names
|
||||
|
||||
|
||||
def _is_environ(node: ast.expr, os_names: set[str], environ_names: set[str]) -> bool:
|
||||
return (
|
||||
isinstance(node, ast.Attribute)
|
||||
and isinstance(node.value, ast.Name)
|
||||
and node.value.id in os_names
|
||||
and node.attr == "environ"
|
||||
) or (isinstance(node, ast.Name) and node.id in environ_names)
|
||||
|
||||
|
||||
def _environment_key_expressions(
|
||||
node: ast.AST,
|
||||
os_names: set[str],
|
||||
getenv_names: set[str],
|
||||
environ_names: set[str],
|
||||
) -> list[ast.expr]:
|
||||
"""Return key expressions used by one direct environment access."""
|
||||
if isinstance(node, ast.Call) and node.args:
|
||||
func = node.func
|
||||
if (
|
||||
isinstance(func, ast.Attribute)
|
||||
and isinstance(func.value, ast.Name)
|
||||
and func.value.id in os_names
|
||||
and func.attr == "getenv"
|
||||
) or (isinstance(func, ast.Name) and func.id in getenv_names):
|
||||
return [node.args[0]]
|
||||
if (
|
||||
isinstance(func, ast.Attribute)
|
||||
and _is_environ(func.value, os_names, environ_names)
|
||||
and func.attr in {"get", "pop", "setdefault"}
|
||||
):
|
||||
return [node.args[0]]
|
||||
elif isinstance(node, ast.Subscript) and _is_environ(node.value, os_names, environ_names):
|
||||
return [node.slice]
|
||||
elif isinstance(node, ast.Compare):
|
||||
expressions: list[ast.expr] = []
|
||||
left_values = [node.left, *node.comparators[:-1]]
|
||||
for left, operator, right in zip(left_values, node.ops, node.comparators):
|
||||
if isinstance(operator, (ast.In, ast.NotIn)) and _is_environ(right, os_names, environ_names):
|
||||
expressions.append(left)
|
||||
return expressions
|
||||
return []
|
||||
|
||||
|
||||
def _resolve_environment_name(expression: ast.expr, constants: dict[str, str]) -> str | None:
|
||||
value: str | None = None
|
||||
if isinstance(expression, ast.Constant) and isinstance(expression.value, str):
|
||||
value = expression.value
|
||||
elif isinstance(expression, ast.Name):
|
||||
value = constants.get(expression.id)
|
||||
return value or None
|
||||
|
||||
|
||||
def _environment_wrapper_parameters(
|
||||
tree: ast.Module,
|
||||
os_names: set[str],
|
||||
getenv_names: set[str],
|
||||
environ_names: set[str],
|
||||
) -> dict[str, set[tuple[str, int | None]]]:
|
||||
"""Find local helpers that forward one parameter as an environment key."""
|
||||
wrappers: dict[str, set[tuple[str, int | None]]] = {}
|
||||
for function in (node for node in tree.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))):
|
||||
positional = [*function.args.posonlyargs, *function.args.args]
|
||||
parameter_indexes: dict[str, int | None] = {parameter.arg: index for index, parameter in enumerate(positional)}
|
||||
parameter_indexes.update({parameter.arg: None for parameter in function.args.kwonlyargs})
|
||||
forwarded: set[tuple[str, int | None]] = set()
|
||||
for node in ast.walk(function):
|
||||
for expression in _environment_key_expressions(node, os_names, getenv_names, environ_names):
|
||||
if isinstance(expression, ast.Name) and expression.id in parameter_indexes:
|
||||
forwarded.add((expression.id, parameter_indexes[expression.id]))
|
||||
if forwarded:
|
||||
wrappers[function.name] = forwarded
|
||||
return wrappers
|
||||
|
||||
|
||||
def _environment_accesses(path: Path) -> set[str]:
|
||||
"""Return statically resolvable names accessed through ``os``."""
|
||||
tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
|
||||
constants = _module_string_constants(tree)
|
||||
os_names, getenv_names, environ_names = _os_import_aliases(tree)
|
||||
wrappers = _environment_wrapper_parameters(tree, os_names, getenv_names, environ_names)
|
||||
names: set[str] = set()
|
||||
|
||||
for node in ast.walk(tree):
|
||||
for expression in _environment_key_expressions(node, os_names, getenv_names, environ_names):
|
||||
if name := _resolve_environment_name(expression, constants):
|
||||
names.add(name)
|
||||
|
||||
if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
|
||||
for parameter_name, position in wrappers.get(node.func.id, set()):
|
||||
argument = node.args[position] if position is not None and position < len(node.args) else None
|
||||
if argument is None:
|
||||
argument = next(
|
||||
(keyword.value for keyword in node.keywords if keyword.arg == parameter_name),
|
||||
None,
|
||||
)
|
||||
if argument is not None and (name := _resolve_environment_name(argument, constants)):
|
||||
names.add(name)
|
||||
|
||||
return names
|
||||
|
||||
|
||||
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.INHERITED_VLLM: 20,
|
||||
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
|
||||
EnvironmentVariableCategory.MODEL_SPECIFIC: 56,
|
||||
EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
|
||||
EnvironmentVariableCategory.INTERNAL: 2,
|
||||
}
|
||||
|
||||
disposition_counts = Counter(
|
||||
item.model_disposition
|
||||
for item in ENVIRONMENT_VARIABLE_INVENTORY.values()
|
||||
if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
|
||||
)
|
||||
assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
|
||||
ModelEnvironmentVariableDisposition.PROMOTE: 33,
|
||||
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
|
||||
ModelEnvironmentVariableDisposition.EXTERNAL: 0,
|
||||
ModelEnvironmentVariableDisposition.INTERNALIZE: 11,
|
||||
ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 6,
|
||||
}
|
||||
|
||||
|
||||
def test_new_public_omni_names_use_project_prefix():
|
||||
"""Grandfather legacy names without allowing more prefix exceptions."""
|
||||
legacy_public_names = {
|
||||
"DIFFUSION_ATTENTION_BACKEND",
|
||||
"DIFFUSION_CACHE_ADAPTER",
|
||||
"DIFFUSION_CACHE_BACKEND",
|
||||
"OMNI_DIFFUSION_PROMPT_EMBED_CACHE",
|
||||
"OMNI_DIFFUSION_PROMPT_EMBED_CACHE_SIZE",
|
||||
"OMNI_DIFFUSION_SESSION_STATE_MANAGER",
|
||||
"OMNI_DIFFUSION_SESSION_STATE_MANAGER_MAX_SESSIONS",
|
||||
"SPEAKER_MAX_UPLOADED",
|
||||
"SPEAKER_SAMPLES_DIR",
|
||||
"VLLM_VIDEO_ASYNC_CHUNK",
|
||||
"VLLM_VIDEO_AUDIO_DELTA_MODE",
|
||||
}
|
||||
public_names = {
|
||||
name
|
||||
for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items()
|
||||
if item.category is EnvironmentVariableCategory.PUBLIC_OMNI
|
||||
}
|
||||
|
||||
assert {name for name in public_names if not name.startswith("VLLM_OMNI_")} == legacy_public_names
|
||||
|
||||
|
||||
def test_statically_resolvable_environment_accesses_are_classified():
|
||||
discovered: set[str] = set()
|
||||
for path in _PACKAGE_ROOT.rglob("*.py"):
|
||||
discovered.update(_environment_accesses(path))
|
||||
|
||||
assert discovered - ENVIRONMENT_VARIABLE_INVENTORY.keys() == set()
|
||||
|
||||
|
||||
def test_model_and_benchmark_inventory_entries_are_still_referenced():
|
||||
"""Reject transitional rows after their implementation has disappeared."""
|
||||
package_source = "\n".join(
|
||||
path.read_text(encoding="utf-8") for path in _PACKAGE_ROOT.rglob("*.py") if path != _INVENTORY_MODULE
|
||||
)
|
||||
transitional_categories = {
|
||||
EnvironmentVariableCategory.MODEL_SPECIFIC,
|
||||
EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL,
|
||||
}
|
||||
|
||||
stale_names = {
|
||||
name
|
||||
for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items()
|
||||
if item.category in transitional_categories and name not in package_source
|
||||
}
|
||||
|
||||
assert stale_names == set()
|
||||
|
||||
|
||||
def test_environment_scanner_covers_indirection_aliases_membership_and_casing():
|
||||
expected_by_path = {
|
||||
"metrics/definitions.py": {
|
||||
"VLLM_OMNI_BENCH_AUDIO_CHANNELS",
|
||||
"VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE",
|
||||
},
|
||||
"entrypoints/openai/video_stream_envs.py": {
|
||||
"VLLM_VIDEO_ASYNC_CHUNK",
|
||||
"VLLM_VIDEO_AUDIO_DELTA_MODE",
|
||||
},
|
||||
"model_executor/models/mimo_audio/mimo_audio.py": {"model_stage"},
|
||||
"model_executor/models/moss_tts/modeling_moss_tts_local_depth.py": {"MOSS_TTS_DEBUG_STOP"},
|
||||
"distributed/ray_utils/utils.py": {"RAY_RAYLET_PID"},
|
||||
}
|
||||
for relative_path, expected in expected_by_path.items():
|
||||
assert expected <= _environment_accesses(_PACKAGE_ROOT / relative_path)
|
||||
|
||||
|
||||
def test_generated_server_storage_environment_names_are_classified():
|
||||
from vllm_omni.config.server_settings import FileBackend
|
||||
|
||||
generated_names = {
|
||||
f"VLLM_OMNI_SERVER_STORAGE__{field_name.upper()}"
|
||||
for field_name in FileBackend.model_fields
|
||||
if field_name != "type"
|
||||
}
|
||||
assert generated_names <= ENVIRONMENT_VARIABLE_INVENTORY.keys()
|
||||
|
||||
|
||||
def test_public_omni_variables_are_in_the_reference_page():
|
||||
reference = _REFERENCE_PAGE.read_text(encoding="utf-8")
|
||||
missing = {
|
||||
name
|
||||
for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items()
|
||||
if item.is_public_omni and f"`{name}`" not in reference
|
||||
}
|
||||
assert missing == set()
|
||||
|
||||
|
||||
def test_secret_values_are_marked_for_redaction():
|
||||
assert {name for name, item in ENVIRONMENT_VARIABLE_INVENTORY.items() if item.redact_value} == {
|
||||
"HF_TOKEN",
|
||||
"HUGGINGFACE_HUB_TOKEN",
|
||||
"OPENAI_API_KEY",
|
||||
}
|
||||
|
||||
|
||||
def test_collect_env_reports_only_safe_public_omni_values(monkeypatch):
|
||||
from collect_env import get_env_vars
|
||||
|
||||
monkeypatch.setenv("DIFFUSION_CACHE_BACKEND", "tea_cache")
|
||||
monkeypatch.setenv("HF_TOKEN", "hf-secret")
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "openai-secret")
|
||||
monkeypatch.setenv("VLLM_OMNI_REPLICA_ID", "internal-replica")
|
||||
|
||||
report = get_env_vars()
|
||||
|
||||
assert "DIFFUSION_CACHE_BACKEND=tea_cache" in report.splitlines()
|
||||
assert "hf-secret" not in report
|
||||
assert "openai-secret" not in report
|
||||
assert "VLLM_OMNI_REPLICA_ID" not in report
|
||||
|
||||
|
||||
@pytest.mark.parametrize("error_type", [ImportError, OSError, AssertionError, RuntimeError])
|
||||
def test_collect_env_survives_inventory_import_error(monkeypatch, error_type):
|
||||
from collect_env import get_env_vars
|
||||
|
||||
real_import = builtins.__import__
|
||||
|
||||
def fail_inventory_import(name, *args, **kwargs):
|
||||
if name == "vllm_omni.config.environment_variable_inventory":
|
||||
raise error_type("broken vllm-omni installation")
|
||||
return real_import(name, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(builtins, "__import__", fail_inventory_import)
|
||||
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0")
|
||||
|
||||
assert "CUDA_VISIBLE_DEVICES=0" in get_env_vars().splitlines()
|
||||
@@ -1,15 +1,24 @@
|
||||
# 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 additive structured Omni config."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import fields
|
||||
from inspect import Parameter, signature
|
||||
from pathlib import Path
|
||||
|
||||
import msgspec
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
from pydantic.fields import FieldInfo
|
||||
from transformers import Qwen3OmniMoeConfig
|
||||
from vllm.config import CacheConfig as VllmCacheConfig
|
||||
from vllm.config import CompilationConfig as VllmCompilationConfig
|
||||
from vllm.config import LoadConfig as VllmLoadConfig
|
||||
from vllm.config import ParallelConfig as VllmParallelConfig
|
||||
from vllm.config import ProfilerConfig as VllmProfilerConfig
|
||||
from vllm.config import SchedulerConfig as VllmSchedulerConfig
|
||||
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
from vllm_omni.config import omni_config as omni_config_module
|
||||
@@ -40,6 +49,7 @@ from vllm_omni.config.stage_config import (
|
||||
merge_pipeline_deploy,
|
||||
)
|
||||
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
|
||||
from vllm_omni.engine.stage_engine_startup import _serialize_stage_config
|
||||
from vllm_omni.engine.stage_init_utils import build_legacy_engine_args_dict
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -82,10 +92,23 @@ def _from_pipeline_key(
|
||||
)
|
||||
|
||||
|
||||
def test_minimax_h3_text_encoder_tp_targets_only_structured_stage_zero() -> None:
|
||||
config = _from_pipeline_key(
|
||||
"minimax_h3_disaggregated",
|
||||
cli_overrides={"text_encoder_tp_size": 4},
|
||||
)
|
||||
|
||||
assert config.stage_by_id(0).parallel_config.tensor_parallel_size == 4
|
||||
diffusion = config.stage_by_id(1)
|
||||
assert diffusion.parallel_config.tensor_parallel_size == 1
|
||||
assert diffusion.parallel_config.text_encoder_tp_size == 1
|
||||
|
||||
|
||||
def test_duplex_session_capacity_propagates_to_every_structured_stage() -> None:
|
||||
omni_config = _from_pipeline_key("personaplex")
|
||||
|
||||
assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [2, 2]
|
||||
assert [stage.model_config.session_mode for stage in omni_config.stage_configs] == ["duplex", "duplex"]
|
||||
|
||||
|
||||
def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) -> None:
|
||||
@@ -95,6 +118,7 @@ def test_non_duplex_deploy_keeps_model_session_capacity_at_one(tmp_path: Path) -
|
||||
omni_config = _from_pipeline_key("personaplex", deploy_config_path=str(deploy_path))
|
||||
|
||||
assert [stage.model_config.duplex_max_sessions for stage in omni_config.stage_configs] == [1, 1]
|
||||
assert [stage.model_config.session_mode for stage in omni_config.stage_configs] == ["turn", "turn"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES))
|
||||
@@ -122,6 +146,7 @@ def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(mod
|
||||
|
||||
engine_args = legacy_stage.yaml_engine_args
|
||||
assert omni_stage.model_config.duplex_max_sessions == engine_args.get("duplex_max_sessions", 1)
|
||||
assert omni_stage.model_config.session_mode == engine_args.get("session_mode", "turn")
|
||||
assert omni_stage.model_config.enforce_eager == engine_args.get("enforce_eager", False)
|
||||
assert omni_stage.load_config.load_format == engine_args.get("load_format", "auto")
|
||||
assert omni_stage.load_config.tokenizer_mode == engine_args.get("tokenizer_mode", "auto")
|
||||
@@ -234,6 +259,36 @@ def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras
|
||||
VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("engine_extras", "cli_overrides"),
|
||||
[
|
||||
({"requires_full_payload_input": False}, {}),
|
||||
({}, {"stage_1_requires_full_payload_input": False}),
|
||||
],
|
||||
ids=["deploy", "stage-cli"],
|
||||
)
|
||||
def test_from_pipeline_config_rejects_full_payload_input_capability_overrides(
|
||||
engine_extras,
|
||||
cli_overrides,
|
||||
):
|
||||
pipeline = _resolve_pipeline_or_skip("qwen3_tts")
|
||||
deploy = DeployConfig(
|
||||
stages=[
|
||||
StageDeployConfig(
|
||||
stage_id=1,
|
||||
engine_extras=engine_extras,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match=r"no structured config owner: requires_full_payload_input"):
|
||||
VllmOmniConfig.from_pipeline_config(
|
||||
pipeline,
|
||||
user_deploy_config=deploy,
|
||||
cli_overrides=cli_overrides,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("cli_overrides", "stage_id"),
|
||||
[
|
||||
@@ -243,9 +298,29 @@ def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras
|
||||
],
|
||||
ids=["global", "ar-stage", "generation-stage"],
|
||||
)
|
||||
def test_from_pipeline_config_rejects_diffusion_only_cli_fields_for_llm_stages(cli_overrides, stage_id):
|
||||
with pytest.raises(ValueError, match=rf"Stage {stage_id} .*no structured config owner: kv_cache_dtype"):
|
||||
_from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides)
|
||||
def test_from_pipeline_config_accepts_upstream_cache_dtype_for_llm_stages(cli_overrides, stage_id):
|
||||
stage = _from_pipeline_key("qwen3_tts", cli_overrides=cli_overrides).stage_by_id(stage_id)
|
||||
|
||||
assert stage.cache_config.cache_dtype == "fp8"
|
||||
assert "cache_dtype" in stage.cache_config._omni_explicit_fields
|
||||
|
||||
|
||||
@pytest.mark.parametrize("field_name", ["data_parallel_master_ip", "data_parallel_address"])
|
||||
def test_from_pipeline_config_normalizes_nested_parallel_config_aliases(field_name):
|
||||
pipeline = _resolve_pipeline_or_skip("qwen3_tts")
|
||||
deploy = DeployConfig(
|
||||
stages=[
|
||||
StageDeployConfig(
|
||||
stage_id=0,
|
||||
engine_extras={"parallel_config": {field_name: "10.0.0.1"}},
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
stage = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy).stage_by_id(0)
|
||||
|
||||
assert stage.parallel_config.data_parallel_master_ip == "10.0.0.1"
|
||||
assert "data_parallel_master_ip" in stage.parallel_config._omni_explicit_fields
|
||||
|
||||
|
||||
def test_from_pipeline_config_accepts_diffusion_only_cli_fields_for_diffusion_stage():
|
||||
@@ -413,10 +488,12 @@ def test_from_pipeline_config_dispatches_async_chunk_processors_without_mutating
|
||||
async_config = _from_pipeline_key("qwen3_tts")
|
||||
assert async_config.stage_by_id(0).custom_process_next_stage_input_func.endswith("talker2code2wav_async_chunk")
|
||||
assert async_config.stage_by_id(1).custom_process_input_func is None
|
||||
assert async_config.stage_by_id(1).model_config.requires_full_payload_input is True
|
||||
|
||||
sync_config = _from_pipeline_key("qwen3_tts", cli_overrides={"async_chunk": False})
|
||||
assert sync_config.stage_by_id(0).custom_process_next_stage_input_func.endswith("talker2code2wav_full_payload")
|
||||
assert sync_config.stage_by_id(1).custom_process_input_func.endswith("talker2code2wav_token_only")
|
||||
assert sync_config.stage_by_id(1).model_config.requires_full_payload_input is True
|
||||
|
||||
assert pipeline.get_stage(0).custom_process_next_stage_input_func.endswith("talker2code2wav_full_payload")
|
||||
assert pipeline.get_stage(1).custom_process_input_func is None
|
||||
@@ -437,6 +514,8 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations():
|
||||
"connector_config",
|
||||
"runtime_config",
|
||||
"parallel_config",
|
||||
"compilation_config",
|
||||
"profiler_config",
|
||||
"quantization_config",
|
||||
}
|
||||
assert "diffusion_config" not in public_fields
|
||||
@@ -455,10 +534,47 @@ def test_runtime_config_fields_match_structured_runtime_scope():
|
||||
"num_gpus",
|
||||
"log_level",
|
||||
"log_stats",
|
||||
"profiler_config",
|
||||
}
|
||||
|
||||
|
||||
def test_upstream_config_compatible_fields_materialize_mapping_inputs():
|
||||
compilation_config = {"backend": "eager"}
|
||||
profiler_config = {"profiler": "cuda"}
|
||||
|
||||
stage_config = VllmOmniARStageConfig(
|
||||
stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
|
||||
compilation_config=compilation_config,
|
||||
profiler_config=profiler_config,
|
||||
)
|
||||
|
||||
assert isinstance(stage_config.compilation_config, VllmCompilationConfig)
|
||||
assert stage_config.compilation_config.backend == "eager"
|
||||
assert isinstance(stage_config.profiler_config, VllmProfilerConfig)
|
||||
assert stage_config.profiler_config.profiler == "cuda"
|
||||
|
||||
|
||||
def test_upstream_config_compatible_fields_validate_mapping_inputs():
|
||||
with pytest.raises(ValidationError):
|
||||
VllmOmniARStageConfig(
|
||||
stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
|
||||
profiler_config={"delay_iterations": -1},
|
||||
)
|
||||
|
||||
|
||||
def test_upstream_config_compatible_fields_keep_prebuilt_config_types():
|
||||
compilation_config = VllmCompilationConfig(backend="eager")
|
||||
profiler_config = VllmProfilerConfig(profiler="cuda")
|
||||
|
||||
stage_config = VllmOmniARStageConfig(
|
||||
stage_pipeline_config=_resolve_pipeline_or_skip("qwen3_tts").stages[0],
|
||||
compilation_config=compilation_config,
|
||||
profiler_config=profiler_config,
|
||||
)
|
||||
|
||||
assert stage_config.compilation_config is compilation_config
|
||||
assert stage_config.profiler_config is profiler_config
|
||||
|
||||
|
||||
def test_sub_config_fields_match_structured_scopes():
|
||||
assert {f.name for f in fields(OmniStageModelConfig)} == {
|
||||
"model",
|
||||
@@ -478,10 +594,12 @@ def test_sub_config_fields_match_structured_scopes():
|
||||
"interleave_mm_strings",
|
||||
"media_io_kwargs",
|
||||
"active_stream_window",
|
||||
"session_mode",
|
||||
"duplex_max_sessions",
|
||||
"enable_sleep_mode",
|
||||
"default_sampling_params",
|
||||
"subtalker_sampling_params",
|
||||
"silence_ban_frames",
|
||||
"has_sampling_extra_args",
|
||||
"custom_voice_dir",
|
||||
"task_type",
|
||||
@@ -489,7 +607,6 @@ def test_sub_config_fields_match_structured_scopes():
|
||||
"enforce_eager",
|
||||
"max_cudagraph_capture_size",
|
||||
"enable_flashinfer_autotune",
|
||||
"compilation_config",
|
||||
"enable_multithread_weight_load",
|
||||
"num_weight_load_threads",
|
||||
"disable_autocast",
|
||||
@@ -498,30 +615,27 @@ def test_sub_config_fields_match_structured_scopes():
|
||||
# legacy engine-args path.
|
||||
"model_subdir",
|
||||
"tokenizer_subdir",
|
||||
"requires_full_payload_input",
|
||||
}
|
||||
assert {f.name for f in fields(OmniStageLoadConfig)} == {
|
||||
vllm_load_fields = {f.name for f in fields(VllmLoadConfig)}
|
||||
assert issubclass(OmniStageLoadConfig, VllmLoadConfig)
|
||||
assert {f.name for f in fields(OmniStageLoadConfig)} == vllm_load_fields | {
|
||||
"tokenizer",
|
||||
"download_dir",
|
||||
"skip_tokenizer_init",
|
||||
"load_format",
|
||||
"tokenizer_mode",
|
||||
"config_format",
|
||||
"skip_mm_profiling",
|
||||
}
|
||||
assert {f.name for f in fields(OmniStageCacheConfig)} == {
|
||||
"kv_cache_memory_bytes",
|
||||
"gpu_memory_utilization",
|
||||
"enable_prefix_caching",
|
||||
assert OmniStageLoadConfig(load_format="PT").load_format == "pt"
|
||||
assert issubclass(OmniStageCacheConfig, VllmCacheConfig)
|
||||
assert {f.name for f in fields(OmniStageCacheConfig)} == {f.name for f in fields(VllmCacheConfig)} | {
|
||||
"disable_hybrid_kv_cache_manager",
|
||||
"mm_processor_cache_gb",
|
||||
"mamba_ssm_cache_dtype",
|
||||
}
|
||||
assert {f.name for f in fields(OmniStageSchedulerConfig)} == {
|
||||
"max_num_seqs",
|
||||
"max_num_batched_tokens",
|
||||
assert issubclass(OmniStageSchedulerConfig, VllmSchedulerConfig)
|
||||
assert {f.name for f in fields(OmniStageSchedulerConfig)} == {f.name for f in fields(VllmSchedulerConfig)} | {
|
||||
"max_model_len",
|
||||
"enable_chunked_prefill",
|
||||
"async_scheduling",
|
||||
}
|
||||
assert {f.name for f in fields(OmniStageConnectorConfig)} == {
|
||||
"async_chunk",
|
||||
@@ -530,44 +644,104 @@ def test_sub_config_fields_match_structured_scopes():
|
||||
"output_connectors",
|
||||
"input_connectors",
|
||||
}
|
||||
assert {f.name for f in fields(OmniStageParallelConfig)} == {
|
||||
"pipeline_parallel_size",
|
||||
"data_parallel_size",
|
||||
"tensor_parallel_size",
|
||||
"enable_expert_parallel",
|
||||
"world_size",
|
||||
}
|
||||
assert {f.name for f in fields(OmniStageDiffusionParallelConfig)} == {
|
||||
"pipeline_parallel_size",
|
||||
"data_parallel_size",
|
||||
"tensor_parallel_size",
|
||||
vllm_parallel_fields = {f.name for f in fields(VllmParallelConfig)}
|
||||
assert issubclass(OmniStageParallelConfig, VllmParallelConfig)
|
||||
assert {f.name for f in fields(OmniStageParallelConfig)} == vllm_parallel_fields
|
||||
assert {f.name for f in fields(OmniStageDiffusionParallelConfig)} == vllm_parallel_fields | {
|
||||
"sequence_parallel_size",
|
||||
"ulysses_degree",
|
||||
"ring_degree",
|
||||
"allgather_degree",
|
||||
"ulysses_mode",
|
||||
"ulysses_a2a_permute",
|
||||
"cfg_parallel_size",
|
||||
"vae_patch_parallel_size",
|
||||
"text_encoder_tp_size",
|
||||
"vae_parallel_mode",
|
||||
"text_encoder_tp_size",
|
||||
"use_hsdp",
|
||||
"mask_sp_padding",
|
||||
"hsdp_shard_size",
|
||||
"hsdp_replicate_size",
|
||||
"enable_expert_parallel",
|
||||
"world_size",
|
||||
}
|
||||
|
||||
|
||||
def test_inherited_sub_configs_initialize_transport_safe_derived_fields():
|
||||
scheduler_config = OmniStageSchedulerConfig(max_num_batched_tokens=4096)
|
||||
assert scheduler_config.max_num_encoder_input_tokens == 4096
|
||||
assert scheduler_config.encoder_cache_size == 4096
|
||||
|
||||
unresolved_scheduler_config = OmniStageSchedulerConfig()
|
||||
assert unresolved_scheduler_config.max_num_encoder_input_tokens is None
|
||||
assert unresolved_scheduler_config.encoder_cache_size is None
|
||||
|
||||
parallel_config = OmniStageParallelConfig(data_parallel_size=3, data_parallel_rank=2)
|
||||
assert parallel_config.data_parallel_index == 2
|
||||
|
||||
diffusion_parallel_config = OmniStageDiffusionParallelConfig(data_parallel_size=3, data_parallel_rank=2)
|
||||
assert diffusion_parallel_config.data_parallel_index == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"config_cls",
|
||||
[
|
||||
OmniStageLoadConfig,
|
||||
OmniStageCacheConfig,
|
||||
OmniStageSchedulerConfig,
|
||||
OmniStageParallelConfig,
|
||||
OmniStageDiffusionParallelConfig,
|
||||
],
|
||||
)
|
||||
def test_inherited_sub_configs_are_keyword_only(config_cls):
|
||||
assert all(parameter.kind is Parameter.KEYWORD_ONLY for parameter in signature(config_cls).parameters.values())
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
config_cls(1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"config_cls",
|
||||
[
|
||||
OmniStageLoadConfig,
|
||||
OmniStageCacheConfig,
|
||||
OmniStageSchedulerConfig,
|
||||
OmniStageParallelConfig,
|
||||
OmniStageDiffusionParallelConfig,
|
||||
],
|
||||
)
|
||||
def test_inherited_sub_configs_remain_msgpack_transport_safe(config_cls):
|
||||
config = config_cls()
|
||||
|
||||
uninitialized_fields = [
|
||||
config_field.name
|
||||
for config_field in fields(config)
|
||||
if isinstance(getattr(config, config_field.name), FieldInfo)
|
||||
]
|
||||
assert uninitialized_fields == []
|
||||
|
||||
msgspec.msgpack.encode(_serialize_stage_config(config))
|
||||
|
||||
|
||||
def test_structured_llm_stage_registration_payloads_remain_msgpack_transport_safe():
|
||||
omni_config = _from_pipeline_key("qwen3_tts")
|
||||
|
||||
for stage_config in omni_config.stage_configs:
|
||||
payload = {
|
||||
"stage_id": stage_config.stage_id,
|
||||
"stage_config": _serialize_stage_config(stage_config),
|
||||
}
|
||||
msgspec.msgpack.encode(payload)
|
||||
|
||||
|
||||
def test_diffusion_parallel_config_fields_cover_legacy_surface():
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
|
||||
legacy_fields = {f.name for f in fields(DiffusionParallelConfig)}
|
||||
structured_fields = {f.name for f in fields(OmniStageDiffusionParallelConfig)}
|
||||
expected_upstream_fields = {"mask_sp_padding"}
|
||||
vllm_parallel_fields = {f.name for f in fields(VllmParallelConfig)}
|
||||
|
||||
assert legacy_fields | expected_upstream_fields <= structured_fields
|
||||
assert structured_fields - legacy_fields - expected_upstream_fields == {"world_size"}
|
||||
assert structured_fields - legacy_fields - expected_upstream_fields - vllm_parallel_fields == set()
|
||||
|
||||
|
||||
def test_diffusion_parallel_config_keeps_current_diffusion_parallel_surface():
|
||||
@@ -610,8 +784,7 @@ def test_diffusion_parallel_config_matches_diffusion_parallel_world_size_for_vae
|
||||
|
||||
def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding():
|
||||
cfg = OmniStageDiffusionParallelConfig(
|
||||
pipeline_parallel_size=2,
|
||||
ulysses_degree=2,
|
||||
ulysses_degree=4,
|
||||
use_hsdp=True,
|
||||
hsdp_shard_size=-1,
|
||||
hsdp_replicate_size=2,
|
||||
@@ -622,10 +795,10 @@ def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding():
|
||||
|
||||
|
||||
def test_diffusion_parallel_config_rejects_hsdp_with_tp_or_dp():
|
||||
with pytest.raises(ValueError, match="cannot be used with TP or DP"):
|
||||
with pytest.raises(ValueError, match="not compatible with TP"):
|
||||
OmniStageDiffusionParallelConfig(tensor_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
|
||||
|
||||
with pytest.raises(ValueError, match="cannot be used with TP or DP"):
|
||||
with pytest.raises(ValueError, match="not compatible with DP"):
|
||||
OmniStageDiffusionParallelConfig(data_parallel_size=2, use_hsdp=True, hsdp_shard_size=2)
|
||||
|
||||
|
||||
@@ -686,6 +859,23 @@ def test_from_pipeline_config_derives_sequence_parallel_size_from_allgather_degr
|
||||
assert stage.parallel_config.world_size == 2
|
||||
|
||||
|
||||
def test_from_pipeline_config_preserves_deployed_ulysses_a2a_permute() -> None:
|
||||
pipeline = _resolve_pipeline_or_skip("dreamzero")
|
||||
deploy = DeployConfig(
|
||||
stages=[
|
||||
StageDeployConfig(
|
||||
stage_id=0,
|
||||
ulysses_a2a_permute=True,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
stage = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy).stage_by_id(0)
|
||||
|
||||
assert isinstance(stage, VllmOmniDiffusionStageConfig)
|
||||
assert stage.parallel_config.ulysses_a2a_permute is True
|
||||
|
||||
|
||||
def test_diffusion_parallel_config_accepts_four_way_guidance_parallelism():
|
||||
cfg = OmniStageDiffusionParallelConfig(cfg_parallel_size=4)
|
||||
|
||||
@@ -780,6 +970,35 @@ def test_structured_diffusion_config_rejects_non_boolean_compile_dynamic():
|
||||
omni_config_module._DiffusionConfigProjection(diffusion_compile_dynamic="false")
|
||||
|
||||
|
||||
def test_from_pipeline_config_routes_ltx2_conv_vae_extra(tmp_path):
|
||||
deploy_path = tmp_path / "dreamzero_ltx2_extras.yaml"
|
||||
deploy_path.write_text(
|
||||
"\n".join(
|
||||
[
|
||||
"pipeline: dreamzero",
|
||||
"async_chunk: false",
|
||||
"stages:",
|
||||
" - stage_id: 0",
|
||||
" extras:",
|
||||
" ltx2_use_conv_vae: true",
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0)
|
||||
|
||||
assert stage.diffusion_config.extras["ltx2_use_conv_vae"] is True
|
||||
|
||||
|
||||
def test_stage_override_routes_ltx2_conv_vae_extra():
|
||||
stage = _from_pipeline_key(
|
||||
"dreamzero",
|
||||
cli_overrides={"stage_0_extras": {"ltx2_use_conv_vae": True}},
|
||||
).stage_by_id(0)
|
||||
|
||||
assert stage.diffusion_config.extras["ltx2_use_conv_vae"] is True
|
||||
|
||||
|
||||
def test_structured_diffusion_config_rejects_invalid_compile_granularity():
|
||||
with pytest.raises(ValidationError, match="diffusion_compile_granularity"):
|
||||
omni_config_module._DiffusionConfigProjection(diffusion_compile_granularity="block")
|
||||
@@ -967,7 +1186,10 @@ def test_diffusion_config_preserves_existing_coercion_hooks():
|
||||
|
||||
|
||||
def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
monkeypatch.setenv("DIFFUSION_CACHE_BACKEND", "TEA_CACHE")
|
||||
monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
|
||||
|
||||
cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
|
||||
diffusion_attention_backend="flash_attn",
|
||||
@@ -977,6 +1199,7 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
|
||||
kv_cache_skip_layers=[2],
|
||||
static_lora_scale=0.25,
|
||||
diffusion_kv_mode="paged_scheduler",
|
||||
diffusion_kv_max_rows_per_request=2,
|
||||
diffusers_load_kwargs=None,
|
||||
diffusers_call_kwargs=None,
|
||||
)
|
||||
@@ -993,7 +1216,10 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
|
||||
assert cfg.diffusers_call_kwargs == {}
|
||||
|
||||
|
||||
def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path):
|
||||
def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path, monkeypatch):
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
|
||||
deploy_path = tmp_path / "dreamzero_diffusion_aliases.yaml"
|
||||
deploy_path.write_text(
|
||||
"\n".join(
|
||||
@@ -1005,6 +1231,7 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
|
||||
" diffusion_attention_backend: flash_attn",
|
||||
" fa_deterministic: true",
|
||||
" diffusion_kv_mode: paged_scheduler",
|
||||
" diffusion_kv_max_rows_per_request: 2",
|
||||
]
|
||||
)
|
||||
)
|
||||
@@ -1018,6 +1245,36 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
|
||||
assert stage.diffusion_config.diffusion_attention_config.default.backend == "flash_attn"
|
||||
assert stage.diffusion_config.fa_deterministic is True
|
||||
assert stage.diffusion_config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
|
||||
assert stage.diffusion_config.diffusion_kv_max_rows_per_request == 2
|
||||
|
||||
|
||||
def test_from_pipeline_config_forwards_fastvideo_vsa_topk(tmp_path, monkeypatch):
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
|
||||
deploy_path = tmp_path / "dreamzero_fastvideo_vsa.yaml"
|
||||
deploy_path.write_text(
|
||||
"\n".join(
|
||||
[
|
||||
"pipeline: dreamzero",
|
||||
"async_chunk: false",
|
||||
"stages:",
|
||||
" - stage_id: 0",
|
||||
" diffusion_attention_backend: FASTVIDEO_VSA",
|
||||
" fastvideo_vsa_topk: 96",
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
stage = _from_pipeline_key(
|
||||
"dreamzero",
|
||||
deploy_config_path=str(deploy_path),
|
||||
).stage_by_id(0)
|
||||
|
||||
attention_config = stage.diffusion_config.diffusion_attention_config
|
||||
assert attention_config.default is not None
|
||||
assert attention_config.default.backend == "FASTVIDEO_VSA"
|
||||
assert attention_config.default.fastvideo_vsa_topk == 96
|
||||
|
||||
|
||||
def test_from_pipeline_config_rejects_reserved_diffusion_kv_mode(tmp_path):
|
||||
@@ -1060,6 +1317,7 @@ def test_diffusion_config_field_classification_covers_current_fields():
|
||||
"prompt_embed_cache_size",
|
||||
"diffusion_kv_cache_dtype",
|
||||
"diffusion_kv_mode",
|
||||
"diffusion_kv_max_rows_per_request",
|
||||
} <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS
|
||||
assert {
|
||||
"revision",
|
||||
@@ -1078,3 +1336,21 @@ def test_diffusion_config_projection_keeps_mapping_quantization_config_serializa
|
||||
cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(quantization_config=quantization_config)
|
||||
|
||||
assert cfg.quantization_config == quantization_config
|
||||
|
||||
|
||||
def test_diffusion_quantization_mapping_reaches_terminal_config(monkeypatch):
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
|
||||
quantization_config = {"method": "int8", "activation_scheme": "dynamic"}
|
||||
cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
|
||||
quantization_config=quantization_config,
|
||||
)
|
||||
|
||||
# Exercise the terminal construction performed by the future typed startup
|
||||
# path without probing ports or loading remote model metadata.
|
||||
monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
|
||||
monkeypatch.setattr(OmniDiffusionConfig, "enrich_config", lambda _self: None)
|
||||
cfg.enrich_config()
|
||||
|
||||
assert cfg.quantization_config is not None
|
||||
assert cfg.quantization_config.get_name() == "int8"
|
||||
@@ -66,3 +66,10 @@ def test_register_resolver_requires_model_type(custom_resolver, clean_pipeline_r
|
||||
"""Ensure that registering a custom resolver to OMNI_PIPELINES requires an explicit model_type."""
|
||||
with pytest.raises(ValueError):
|
||||
register_pipeline(custom_resolver)
|
||||
|
||||
|
||||
def test_minimax_h3_disaggregation_is_explicit_opt_in():
|
||||
assert "minimax_h3" not in OMNI_PIPELINES
|
||||
pipeline = OMNI_PIPELINES["minimax_h3_disaggregated"]
|
||||
assert isinstance(pipeline, PipelineConfig)
|
||||
assert pipeline.model_type == "minimax_h3_disaggregated"
|
||||
+8
-47
@@ -1,7 +1,10 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""
|
||||
Root pytest entrypoint for the vLLM-Omni test suite.
|
||||
|
||||
- `tests/conftest.py` stays thin: plugin registration + compatibility re-exports.
|
||||
- `tests/conftest.py` stays thin: plugin registration only.
|
||||
- Importable utilities live under `tests/helpers/`.
|
||||
- Helper unit tests live under `tests/helpers/tests/`.
|
||||
- Fixtures live under `tests/helpers/fixtures/` and are loaded via `pytest_plugins`.
|
||||
@@ -19,57 +22,15 @@ bootstrap_vllm_layer_custom_op_modules()
|
||||
|
||||
pytest_plugins = (
|
||||
"tests.helpers.fixtures.config",
|
||||
"tests.helpers.fixtures.env",
|
||||
"tests.helpers.fixtures.clean",
|
||||
"tests.helpers.fixtures.log",
|
||||
"tests.helpers.fixtures.run_args",
|
||||
"tests.helpers.fixtures.media",
|
||||
"tests.helpers.fixtures.pytest_collection",
|
||||
"tests.helpers.fixtures.pytest_run_args",
|
||||
"tests.helpers.fixtures.runtime",
|
||||
"tests.helpers.fixtures.speaker_cache",
|
||||
)
|
||||
|
||||
|
||||
def pytest_terminal_summary(terminalreporter, exitstatus, config):
|
||||
# Marker for Buildkite log folding before pytest summary lines.
|
||||
terminalreporter.write_line("--- Running Summary")
|
||||
|
||||
|
||||
# Backward-compatible lazy re-exports.
|
||||
# (Many tests still import from `tests.conftest`; migrate these imports to `tests.helpers.*` over time.)
|
||||
# Keep these lazy so conftest import does not trigger heavy helper dependencies.
|
||||
_ASSERTION_EXPORT_NAMES = (
|
||||
"assert_audio_speech_response",
|
||||
"assert_diffusion_response",
|
||||
"assert_http_error",
|
||||
"assert_image_diffusion_response",
|
||||
"assert_image_valid",
|
||||
"assert_omni_response",
|
||||
"assert_video_diffusion_response",
|
||||
"assert_video_valid",
|
||||
)
|
||||
_MEDIA_EXPORT_NAMES = (
|
||||
"convert_audio_bytes_to_text",
|
||||
"convert_audio_file_to_text",
|
||||
"cosine_similarity_text",
|
||||
"decode_b64_image",
|
||||
"generate_synthetic_audio",
|
||||
"generate_synthetic_image",
|
||||
"generate_synthetic_video",
|
||||
)
|
||||
_STAGE_CONFIG_EXPORT_NAMES = ("modify_stage_config",)
|
||||
_RUNTIME_EXPORT_NAMES = (
|
||||
"DiffusionResponse",
|
||||
"HttpResponse",
|
||||
"OmniResponse",
|
||||
"OmniRunner",
|
||||
"OmniRunnerHandler",
|
||||
"OmniServer",
|
||||
"OmniServerParams",
|
||||
"OmniServerStageCli",
|
||||
"OpenAIClientHandler",
|
||||
"dummy_messages_from_mix_data",
|
||||
)
|
||||
_LAZY_EXPORT_MODULES = {
|
||||
**{name: "tests.helpers.assertions" for name in _ASSERTION_EXPORT_NAMES},
|
||||
**{name: "tests.helpers.media" for name in _MEDIA_EXPORT_NAMES},
|
||||
**{name: "tests.helpers.stage_config" for name in _STAGE_CONFIG_EXPORT_NAMES},
|
||||
**{name: "tests.helpers.runtime" for name in _RUNTIME_EXPORT_NAMES},
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
"""Tests for ``OmniARScheduler._drop_aborted_queued_requests``.
|
||||
|
||||
``schedule()`` sweeps ``FINISHED_ABORTED`` requests out of the queues before
|
||||
handing control to upstream, because upstream ``Scheduler.schedule()`` raises
|
||||
``RuntimeError: Invalid request status: FINISHED_ABORTED`` for any request it
|
||||
admits in a finished state -- which kills the stage's engine core, not just
|
||||
the request.
|
||||
|
||||
Regression for the sweep missing ``skipped_waiting``, the third queue
|
||||
upstream admits from. An aborted duplex session parked there was re-selected
|
||||
by ``_select_waiting_queue_for_scheduling`` on a later tick and crashed the
|
||||
stage.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
|
||||
from vllm.v1.request import RequestStatus
|
||||
|
||||
import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
_POLICIES = [
|
||||
pytest.param(SchedulingPolicy.FCFS, id="fcfs"),
|
||||
pytest.param(SchedulingPolicy.PRIORITY, id="priority"),
|
||||
]
|
||||
|
||||
|
||||
class _StubRequest:
|
||||
"""Minimal ``Request`` stub with the surface the sweep exercises.
|
||||
|
||||
Carries the ``(priority, arrival_time)`` ordering ``PriorityRequestQueue``
|
||||
relies on so the same stub works under both scheduling policies.
|
||||
"""
|
||||
|
||||
def __init__(self, request_id: str, status: RequestStatus, arrival_time: float = 0.0) -> None:
|
||||
self.request_id = request_id
|
||||
self.status = status
|
||||
self.priority = 0
|
||||
self.arrival_time = arrival_time
|
||||
|
||||
def __lt__(self, other: _StubRequest) -> bool:
|
||||
return (self.priority, self.arrival_time) < (other.priority, other.arrival_time)
|
||||
|
||||
|
||||
def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), skipped_waiting=(), running=()):
|
||||
"""Bare scheduler with just the surface the sweep reads."""
|
||||
scheduler = OmniARScheduler.__new__(OmniARScheduler)
|
||||
scheduler.waiting = create_request_queue(policy)
|
||||
for req in waiting:
|
||||
scheduler.waiting.add_request(req)
|
||||
scheduler.skipped_waiting = create_request_queue(policy)
|
||||
for req in skipped_waiting:
|
||||
scheduler.skipped_waiting.add_request(req)
|
||||
scheduler.running = list(running)
|
||||
return scheduler
|
||||
|
||||
|
||||
@pytest.mark.parametrize("policy", _POLICIES)
|
||||
def test_schedule_sweeps_skipped_waiting_before_upstream_selection(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
policy: SchedulingPolicy,
|
||||
) -> None:
|
||||
"""The public schedule path must remove aborted skipped requests first."""
|
||||
aborted = _StubRequest("req-aborted", RequestStatus.FINISHED_ABORTED)
|
||||
scheduler = _make_scheduler(policy, skipped_waiting=[aborted])
|
||||
scheduler.chunk_transfer_adapter = None
|
||||
scheduler.input_coordinator = None
|
||||
scheduler.num_waiting_for_streaming_input = 0
|
||||
scheduler._process_pending_omni_inputs = lambda model_mode: None
|
||||
scheduler._resync_streaming_input_counter = lambda: None
|
||||
scheduler._should_defer_waiting_admission = lambda: False
|
||||
scheduler._restore_omni_wait_queues = lambda: None
|
||||
scheduler._postprocess_omni_schedule_output = lambda *args, **kwargs: None
|
||||
scheduler.get_finished_requests_needing_kv_transfer = lambda: {}
|
||||
scheduler._wrap_omni_scheduler_output = lambda output, **kwargs: output
|
||||
|
||||
sentinel = object()
|
||||
|
||||
def upstream_schedule(_self, _throttle_prefills: bool = False):
|
||||
assert list(scheduler.skipped_waiting) == []
|
||||
return sentinel
|
||||
|
||||
monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", upstream_schedule)
|
||||
|
||||
assert OmniARScheduler.schedule(scheduler) is sentinel
|
||||
@@ -100,6 +100,7 @@ def _make_drain_sched(session: Request) -> MagicMock:
|
||||
sched.pending_stop_after_extraction = set()
|
||||
sched.connector = None
|
||||
sched.kv_cache_manager.take_events.return_value = None
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
sched.finished_req_ids_dict = {}
|
||||
sched.make_stats.return_value = None
|
||||
return sched
|
||||
@@ -141,13 +142,16 @@ def test_exact_drain_delivers_new_segment_frame() -> None:
|
||||
session.num_computed_tokens = 6
|
||||
session.num_output_placeholders = 1
|
||||
session.num_in_flight_tokens = 1
|
||||
session.async_tokens_to_discard = 1
|
||||
|
||||
_replace_streaming_session(session)
|
||||
assert session.num_stale_output_tokens == 1
|
||||
assert session.async_tokens_to_discard == 1
|
||||
|
||||
sched = _make_drain_sched(session)
|
||||
assert _run_step(sched, session, num_scheduled=1, token=42) is False # late frame dropped
|
||||
assert session.num_stale_output_tokens == 0
|
||||
assert session.async_tokens_to_discard == 0
|
||||
assert _run_step(sched, session, num_scheduled=1, token=43) is True # new segment survives
|
||||
|
||||
|
||||
|
||||
@@ -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 Omni AR streaming-session async placeholder handling."""
|
||||
|
||||
@@ -18,21 +18,26 @@ from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.outputs import ModelRunnerOutput
|
||||
from vllm.v1.request import Request, RequestStatus, StreamingUpdate
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler
|
||||
|
||||
# isort: on
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_scheduler(*, stage_id: int = 0) -> OmniARScheduler:
|
||||
def _make_scheduler(*, stage_id: int = 0, session_mode: str = "turn") -> OmniARScheduler:
|
||||
sched = OmniARScheduler.__new__(OmniARScheduler)
|
||||
sched._new_prompt_len_snapshot = {}
|
||||
sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(stage_id=stage_id))
|
||||
sched.vllm_config = SimpleNamespace(
|
||||
model_config=SimpleNamespace(stage_id=stage_id, session_mode=session_mode),
|
||||
)
|
||||
sched.num_waiting_for_streaming_input = 0
|
||||
sched.log_stats = False
|
||||
sched.chunk_transfer_adapter = None
|
||||
sched.skipped_waiting = set()
|
||||
sched._free_request_blocks = MagicMock()
|
||||
sched.encoder_cache_manager = MagicMock()
|
||||
sched._inflight_prefills = set()
|
||||
return sched
|
||||
|
||||
|
||||
@@ -61,6 +66,9 @@ def _run_resumable_segment_stop(
|
||||
session: Request,
|
||||
*,
|
||||
session_finished: bool = False,
|
||||
handle_stopped=None,
|
||||
chunk_transfer_adapter=None,
|
||||
inter_stage_output=None,
|
||||
):
|
||||
sched = MagicMock()
|
||||
sched.requests = {session.request_id: session}
|
||||
@@ -72,11 +80,14 @@ def _run_resumable_segment_stop(
|
||||
return [42], True
|
||||
|
||||
sched._update_request_with_output.side_effect = stop_request
|
||||
sched._handle_stopped_request.return_value = session_finished
|
||||
# vLLM 0.26 returns (kv_xfer_params, ec_xfer_params); an unconfigured
|
||||
# MagicMock iterates empty and fails to unpack at the call site.
|
||||
sched._free_request.return_value = (None, None)
|
||||
sched.chunk_transfer_adapter = None
|
||||
if handle_stopped is None:
|
||||
sched._handle_stopped_request.return_value = session_finished
|
||||
else:
|
||||
sched._handle_stopped_request.side_effect = handle_stopped
|
||||
sched.chunk_transfer_adapter = chunk_transfer_adapter
|
||||
sched.running = [session]
|
||||
sched.waiting_for_transfer_free = set()
|
||||
sched.transfer_triggered_requests = set()
|
||||
@@ -84,6 +95,7 @@ def _run_resumable_segment_stop(
|
||||
sched.pending_stop_after_extraction = set()
|
||||
sched.connector = None
|
||||
sched.kv_cache_manager.take_events.return_value = None
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
sched.finished_req_ids_dict = {}
|
||||
sched.make_stats.return_value = None
|
||||
|
||||
@@ -102,6 +114,7 @@ def _run_resumable_segment_stop(
|
||||
model_runner_output.cudagraph_stats = None
|
||||
model_runner_output.req_id_to_index = {session.request_id: 0}
|
||||
model_runner_output.routed_experts = None
|
||||
model_runner_output.inter_stage_outputs = [inter_stage_output] if inter_stage_output is not None else None
|
||||
|
||||
return OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
|
||||
|
||||
@@ -157,6 +170,44 @@ def test_update_from_output_settles_in_flight_tokens() -> None:
|
||||
assert session.num_in_flight_tokens == 0
|
||||
|
||||
|
||||
def test_resumable_segment_boundary_keeps_pre_transition_send_watermark() -> None:
|
||||
"""The old segment's transfer must not observe the next segment's reset.
|
||||
|
||||
A queued streaming update can replace the same mutable Request while the
|
||||
stop output is still being handled. The connector needs the confirmed
|
||||
token count from before that replacement to avoid dropping the boundary.
|
||||
"""
|
||||
session = _make_request()
|
||||
session.status = RequestStatus.RUNNING
|
||||
session.resumable = True
|
||||
session.num_computed_tokens = 26
|
||||
adapter = MagicMock()
|
||||
adapter._confirmed_num_computed_tokens.return_value = 26
|
||||
inter_stage_output = {"codes": {"audio": [7]}}
|
||||
|
||||
def replace_with_next_segment(request: Request) -> bool:
|
||||
request.num_computed_tokens = 0
|
||||
request._omni_segment_generation = 1
|
||||
request.status = RequestStatus.WAITING
|
||||
return False
|
||||
|
||||
_run_resumable_segment_stop(
|
||||
session,
|
||||
handle_stopped=replace_with_next_segment,
|
||||
chunk_transfer_adapter=adapter,
|
||||
inter_stage_output=inter_stage_output,
|
||||
)
|
||||
|
||||
adapter.save_async.assert_called_once_with(
|
||||
inter_stage_output,
|
||||
session,
|
||||
True,
|
||||
new_token_ids=[42],
|
||||
confirmed_num_computed_tokens=26,
|
||||
segment_generation=0,
|
||||
)
|
||||
|
||||
|
||||
def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> None:
|
||||
"""A decode step that neither stops nor carries an inter-stage payload.
|
||||
|
||||
@@ -181,6 +232,7 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
|
||||
sched.pending_stop_after_extraction = set()
|
||||
sched.connector = None
|
||||
sched.kv_cache_manager.take_events.return_value = None
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
sched.finished_req_ids_dict = {}
|
||||
sched.make_stats.return_value = None
|
||||
|
||||
@@ -207,6 +259,171 @@ def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> Non
|
||||
sched.chunk_transfer_adapter.save_async.assert_not_called()
|
||||
|
||||
|
||||
def test_queued_streaming_update_on_async_stop_fences_in_flight_once() -> None:
|
||||
"""Native duplex + async: the next unit is usually already queued when the
|
||||
current one stops, so _handle_stopped_request applies the update and the
|
||||
stop site below it fences the same in-flight decode a second time. An
|
||||
accumulated residue outlives the drain and swallows the next unit's
|
||||
listen/speak, which the client sees as silence until its timeout.
|
||||
"""
|
||||
session = _make_request()
|
||||
session.status = RequestStatus.RUNNING
|
||||
session.resumable = True
|
||||
session.append_output_token_ids([7])
|
||||
session.num_computed_tokens = 4
|
||||
session.num_output_placeholders = 1
|
||||
# This step's frame (settled to 0 by update_from_output) plus one extra
|
||||
# async decode still in flight.
|
||||
session.num_in_flight_tokens = 2
|
||||
|
||||
sched = _make_scheduler(stage_id=0)
|
||||
sched._enqueue_waiting_request = MagicMock()
|
||||
|
||||
def handle_stopped(request: Request) -> bool:
|
||||
sched._update_request_as_session(request, _make_update([10, 20]))
|
||||
return False
|
||||
|
||||
_run_resumable_segment_stop(session, handle_stopped=handle_stopped)
|
||||
|
||||
# Exactly the one unreported decode, so the drain reaches zero before the
|
||||
# new segment's first frame arrives.
|
||||
assert session.num_stale_output_tokens == session.num_in_flight_tokens == 1
|
||||
|
||||
|
||||
def test_stale_async_frame_is_dropped_before_output_processing() -> None:
|
||||
session = _make_request()
|
||||
session.status = RequestStatus.RUNNING
|
||||
session.num_in_flight_tokens = 2
|
||||
session.num_computed_tokens = session.num_tokens + 1
|
||||
session.num_output_placeholders = 1
|
||||
session.async_tokens_to_discard = 1
|
||||
session.sampling_params = SimpleNamespace(num_logprobs=1)
|
||||
num_computed_tokens = session.num_computed_tokens
|
||||
num_output_placeholders = session.num_output_placeholders
|
||||
|
||||
sched = MagicMock()
|
||||
sched.requests = {session.request_id: session}
|
||||
sched.perf_metrics = None
|
||||
sched.structured_output_manager.should_advance.return_value = False
|
||||
|
||||
def discard_stale_output(request: Request, token_ids: list[int]) -> tuple[list[int], bool]:
|
||||
request.async_tokens_to_discard = 0
|
||||
return token_ids, False
|
||||
|
||||
sched._update_request_with_output.side_effect = discard_stale_output
|
||||
sched._process_kv_transfer_trigger.return_value = False
|
||||
sched.chunk_transfer_adapter = MagicMock()
|
||||
sched.running = [session]
|
||||
sched.waiting_for_transfer_free = set()
|
||||
sched.transfer_triggered_requests = set()
|
||||
sched.active_kv_transfers = set()
|
||||
sched.pending_stop_after_extraction = set()
|
||||
sched.connector = None
|
||||
sched.kv_cache_manager.take_events.return_value = None
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
sched.finished_req_ids_dict = {}
|
||||
sched.make_stats.return_value = None
|
||||
|
||||
scheduler_output = MagicMock(spec=SchedulerOutput)
|
||||
scheduler_output.num_scheduled_tokens = {session.request_id: 1}
|
||||
scheduler_output.scheduled_spec_decode_tokens = {}
|
||||
scheduler_output.num_invalid_spec_tokens = 0
|
||||
|
||||
model_runner_output = MagicMock(spec=ModelRunnerOutput)
|
||||
model_runner_output.sampled_token_ids = [[42]]
|
||||
model_runner_output.logprobs = None
|
||||
model_runner_output.prompt_logprobs_dict = {}
|
||||
model_runner_output.pooler_output = None
|
||||
model_runner_output.multimodal_outputs = None
|
||||
model_runner_output.inter_stage_outputs = [{"hidden": object()}]
|
||||
model_runner_output.num_nans_in_logits = None
|
||||
model_runner_output.kv_connector_output = None
|
||||
model_runner_output.cudagraph_stats = None
|
||||
model_runner_output.req_id_to_index = {session.request_id: 0}
|
||||
model_runner_output.routed_experts = None
|
||||
|
||||
OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
|
||||
|
||||
assert session.async_tokens_to_discard == 0
|
||||
assert session.status == RequestStatus.RUNNING
|
||||
assert session.num_computed_tokens == num_computed_tokens
|
||||
assert session.num_output_placeholders == num_output_placeholders
|
||||
sched.chunk_transfer_adapter.save_async.assert_not_called()
|
||||
|
||||
session.sampling_params = SamplingParams(max_tokens=8)
|
||||
next_payload = {"hidden": object()}
|
||||
sched._update_request_with_output.side_effect = None
|
||||
sched._update_request_with_output.return_value = ([43], False)
|
||||
model_runner_output.sampled_token_ids = [[43]]
|
||||
model_runner_output.inter_stage_outputs = [next_payload]
|
||||
|
||||
OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
|
||||
|
||||
sched.chunk_transfer_adapter.save_async.assert_called_once_with(
|
||||
next_payload,
|
||||
session,
|
||||
False,
|
||||
new_token_ids=[43],
|
||||
confirmed_num_computed_tokens=None,
|
||||
)
|
||||
|
||||
|
||||
def test_legacy_stale_async_marker_bypasses_real_placeholder_accounting() -> None:
|
||||
request = _make_request()
|
||||
request.status = RequestStatus.RUNNING
|
||||
request.num_in_flight_tokens = 1
|
||||
request.num_output_placeholders = 0
|
||||
request.async_tokens_to_discard = 1
|
||||
request.num_stale_output_tokens = 0
|
||||
|
||||
sched = OmniARAsyncScheduler.__new__(OmniARAsyncScheduler)
|
||||
sched.requests = {request.request_id: request}
|
||||
sched.perf_metrics = None
|
||||
sched.chunk_transfer_adapter = None
|
||||
sched.connector = None
|
||||
sched.ec_connector = None
|
||||
sched.waiting_for_transfer_free = set()
|
||||
sched.transfer_triggered_requests = set()
|
||||
sched.active_kv_transfers = set()
|
||||
sched.pending_stop_after_extraction = set()
|
||||
sched.finished_req_ids_dict = {}
|
||||
sched._new_prompt_len_snapshot = {}
|
||||
sched.kv_cache_manager = MagicMock()
|
||||
sched.kv_cache_manager.take_events.return_value = None
|
||||
sched._remove_stopped_requests_from_queues = MagicMock()
|
||||
sched._handle_failed_kv_load_outputs = MagicMock(return_value=[])
|
||||
sched._cleanup_kv_tracking = MagicMock()
|
||||
sched._aggregate_kv_connector_stats = MagicMock(return_value=None)
|
||||
sched._publish_kv_cache_events = MagicMock()
|
||||
sched._attach_finished_request_sets = MagicMock()
|
||||
sched._attach_scheduler_stats = MagicMock()
|
||||
sched._capture_omni_connector_output = MagicMock()
|
||||
|
||||
scheduler_output = MagicMock(spec=SchedulerOutput)
|
||||
scheduler_output.num_scheduled_tokens = {request.request_id: 1}
|
||||
scheduler_output.scheduled_spec_decode_tokens = {}
|
||||
scheduler_output.num_invalid_spec_tokens = 0
|
||||
|
||||
model_runner_output = MagicMock(spec=ModelRunnerOutput)
|
||||
model_runner_output.sampled_token_ids = [[42]]
|
||||
model_runner_output.logprobs = None
|
||||
model_runner_output.prompt_logprobs_dict = {}
|
||||
model_runner_output.pooler_output = None
|
||||
model_runner_output.multimodal_outputs = None
|
||||
model_runner_output.inter_stage_outputs = None
|
||||
model_runner_output.num_nans_in_logits = None
|
||||
model_runner_output.kv_connector_output = None
|
||||
model_runner_output.cudagraph_stats = None
|
||||
model_runner_output.req_id_to_index = {request.request_id: 0}
|
||||
model_runner_output.routed_experts = None
|
||||
|
||||
OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
|
||||
|
||||
assert request.async_tokens_to_discard == 0
|
||||
assert request.num_output_placeholders == 0
|
||||
assert request._output_token_ids == []
|
||||
|
||||
|
||||
def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() -> None:
|
||||
sched = _make_scheduler(stage_id=0)
|
||||
session = _make_request()
|
||||
@@ -276,29 +493,33 @@ def test_explicit_streaming_payload_replaces_placeholder_prompt() -> None:
|
||||
"meta": {"turn_eos_token_id": 99},
|
||||
}
|
||||
assert session.status == RequestStatus.WAITING
|
||||
sched._free_request_blocks.assert_called_once_with(session)
|
||||
sched.encoder_cache_manager.free.assert_called_once_with(session)
|
||||
|
||||
|
||||
def test_model_intermediate_streaming_payload_replaces_computed_prompt() -> None:
|
||||
def test_explicit_model_intermediate_prompt_replacement_releases_cache_and_watermark() -> None:
|
||||
sched = _make_scheduler(stage_id=1)
|
||||
session = _make_request()
|
||||
sched.chunk_transfer_adapter = SimpleNamespace(
|
||||
receives_chunks=False,
|
||||
segment_finished_requests=set(),
|
||||
requests_num_chunks_sent={session.external_req_id: 59},
|
||||
)
|
||||
session = _make_request()
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
session.prompt_token_ids = [0] * 59
|
||||
session._all_token_ids.clear()
|
||||
session._all_token_ids.extend(session.prompt_token_ids)
|
||||
session.num_prompt_tokens = 59
|
||||
session.num_computed_tokens = 59
|
||||
session.num_in_flight_tokens = 2
|
||||
update = _make_update([0] * 10)
|
||||
update.additional_information = None
|
||||
update.model_intermediate_buffer = {
|
||||
"ids": {"tts": list(range(8))},
|
||||
"hidden_states": {"tts": [[0.0]] * 8},
|
||||
"meta": {
|
||||
"replace_streaming_prompt": True,
|
||||
"next_stage_prompt_len": 10,
|
||||
"replace_streaming_prompt": True,
|
||||
},
|
||||
}
|
||||
|
||||
@@ -308,6 +529,110 @@ def test_model_intermediate_streaming_payload_replaces_computed_prompt() -> None
|
||||
assert list(session._all_token_ids) == [0] * 10
|
||||
assert session.num_prompt_tokens == 10
|
||||
assert session.num_computed_tokens == 0
|
||||
assert session.num_stale_output_tokens == 2
|
||||
assert session.additional_information is None
|
||||
assert session.model_intermediate_buffer == update.model_intermediate_buffer
|
||||
assert session.status == RequestStatus.WAITING
|
||||
assert sched.chunk_transfer_adapter.requests_num_chunks_sent == {}
|
||||
sched._free_request_blocks.assert_called_once_with(session)
|
||||
sched.encoder_cache_manager.free.assert_called_once_with(session)
|
||||
|
||||
|
||||
def test_ready_async_chunk_prompt_replacement_releases_stale_kv_once() -> None:
|
||||
sched = _make_scheduler(stage_id=1)
|
||||
session = _make_request()
|
||||
session.external_req_id = "external-ar-streaming-test"
|
||||
session.num_in_flight_tokens = 2
|
||||
# _update_request_as_session() may have already fenced this frame before
|
||||
# the connector marks the explicit replacement ready.
|
||||
session.num_stale_output_tokens = 2
|
||||
session.num_output_placeholders = 2
|
||||
session.spec_token_ids = [-1, -1]
|
||||
sched.requests = {session.request_id: session}
|
||||
sched._inflight_prefills.add(session)
|
||||
sched.chunk_transfer_adapter = SimpleNamespace(
|
||||
replaced_streaming_prompt_ids={session.request_id},
|
||||
requests_with_ready_chunks={session.request_id},
|
||||
requests_num_chunks_sent={session.external_req_id: 4090},
|
||||
)
|
||||
|
||||
sched._reset_ready_async_chunk_replacements()
|
||||
sched._reset_ready_async_chunk_replacements()
|
||||
|
||||
sched._free_request_blocks.assert_called_once_with(session)
|
||||
sched.encoder_cache_manager.free.assert_called_once_with(session)
|
||||
assert session not in sched._inflight_prefills
|
||||
assert session.num_stale_output_tokens == 2
|
||||
assert session.num_output_placeholders == 0
|
||||
assert session.spec_token_ids == []
|
||||
assert sched.chunk_transfer_adapter.replaced_streaming_prompt_ids == set()
|
||||
assert sched.chunk_transfer_adapter.requests_with_ready_chunks == {session.request_id}
|
||||
assert sched.chunk_transfer_adapter.requests_num_chunks_sent == {}
|
||||
|
||||
|
||||
def test_chunk_segment_cleanup_keeps_requeued_resumable_receiver() -> None:
|
||||
"""A WAITING_FOR_CHUNK stop must not delete its newly parked session."""
|
||||
session = _make_request()
|
||||
session.resumable = True
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
|
||||
def queue(*requests):
|
||||
result = MagicMock()
|
||||
result.requests = set(requests)
|
||||
result.add_request.side_effect = result.requests.add
|
||||
result.remove_requests.side_effect = result.requests.difference_update
|
||||
return result
|
||||
|
||||
sched = OmniARScheduler.__new__(OmniARScheduler)
|
||||
sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(session_mode="duplex"))
|
||||
sched.running = []
|
||||
sched.waiting = queue()
|
||||
sched.skipped_waiting = queue(session)
|
||||
sched.num_waiting_for_streaming_input = 1
|
||||
sched.chunk_transfer_adapter = SimpleNamespace(
|
||||
receives_chunks=True,
|
||||
segment_finished_requests={session.request_id},
|
||||
)
|
||||
|
||||
sched._resume_downstream_chunk_receiver(session)
|
||||
sched._remove_stopped_requests_from_queues(set(), {session})
|
||||
|
||||
assert session.status == RequestStatus.WAITING
|
||||
assert session in sched.waiting.requests
|
||||
assert session not in sched.skipped_waiting.requests
|
||||
assert sched.num_waiting_for_streaming_input == 0
|
||||
assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("receives_chunks", "session_mode"),
|
||||
[
|
||||
(False, "duplex"),
|
||||
(True, "turn"),
|
||||
],
|
||||
)
|
||||
def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked(
|
||||
receives_chunks: bool,
|
||||
session_mode: str,
|
||||
) -> None:
|
||||
"""Only duplex connector-driven receivers resume without an update."""
|
||||
session = _make_request()
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
|
||||
sched = OmniARScheduler.__new__(OmniARScheduler)
|
||||
sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(session_mode=session_mode))
|
||||
sched.num_waiting_for_streaming_input = 1
|
||||
sched.chunk_transfer_adapter = SimpleNamespace(
|
||||
receives_chunks=receives_chunks,
|
||||
segment_finished_requests={session.request_id},
|
||||
)
|
||||
sched.skipped_waiting = MagicMock()
|
||||
sched._enqueue_waiting_request = MagicMock()
|
||||
|
||||
sched._resume_downstream_chunk_receiver(session)
|
||||
|
||||
assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
assert sched.num_waiting_for_streaming_input == 1
|
||||
assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
|
||||
sched.skipped_waiting.remove_requests.assert_not_called()
|
||||
sched._enqueue_waiting_request.assert_not_called()
|
||||
@@ -76,6 +76,9 @@ def _make_scheduler(
|
||||
else:
|
||||
hf_config.voxcpm2_runtime_config = runtime_config
|
||||
sched.vllm_config = SimpleNamespace(model_config=model_config)
|
||||
# Upstream ``Scheduler.__init__`` always creates ``skipped_waiting``; the
|
||||
# abort sweep at the top of ``schedule()`` reads it.
|
||||
sched.skipped_waiting = _MockQueue()
|
||||
return sched
|
||||
|
||||
|
||||
|
||||
@@ -138,6 +138,7 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
|
||||
sched.recompute_kv_load_failures = False
|
||||
sched.connector = None
|
||||
sched.kv_cache_manager.take_events.return_value = None
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
sched.finished_req_ids_dict = {}
|
||||
sched.make_stats.return_value = None
|
||||
|
||||
@@ -410,12 +411,13 @@ class TestRealignRequestStatusToQueues:
|
||||
assert stale.status == RequestStatus.RUNNING
|
||||
assert clean.status == RequestStatus.RUNNING
|
||||
|
||||
def test_finished_request_is_skipped(self) -> None:
|
||||
def test_non_resumable_finished_request_is_skipped(self) -> None:
|
||||
"""Already-finished requests must not be touched -- they may
|
||||
have legitimate finished statuses (FINISHED_STOPPED etc.) that
|
||||
a status flip would corrupt.
|
||||
"""
|
||||
req = _make_request(request_id="req-finished")
|
||||
req.resumable = False
|
||||
req.status = RequestStatus.FINISHED_STOPPED
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
|
||||
@@ -0,0 +1,180 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""v0.28 deferred KV-block free fence in the Omni schedulers (#6606 review).
|
||||
|
||||
Upstream v0.28 (defer_block_free) advances ``sched_step_seq`` for every
|
||||
non-empty scheduled step and, at the top of ``update_from_output()``,
|
||||
advances ``processed_step_seq`` and drains the deferred frees. The Omni
|
||||
schedulers each carried only half of that contract:
|
||||
|
||||
- ``OmniARScheduler.schedule()`` delegates to ``super().schedule()`` (fence
|
||||
advances) but its reimplemented ``update_from_output()`` never advanced
|
||||
``processed_step_seq`` nor drained -> deferred blocks accumulate without
|
||||
bound under async/PP KV-consumer configurations.
|
||||
- ``OmniGenerationScheduler`` reimplements both: the fast-path
|
||||
``schedule()`` never advanced ``sched_step_seq`` (so finished requests
|
||||
freed blocks an in-flight GPU step could still be writing) and its
|
||||
``update_from_output()`` never drained.
|
||||
|
||||
These tests pin the restored contract at the call boundary; the drain and
|
||||
free logic itself is upstream-tested.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
|
||||
# before Request / RequestStatus are bound in this module.
|
||||
# isort: off
|
||||
import vllm_omni # noqa: F401 - import for side effects (patch vLLM)
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.core.sched.interface import PauseState
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.core.sched.request_queue import SchedulingPolicy
|
||||
from vllm.v1.outputs import ModelRunnerOutput
|
||||
from vllm.v1.request import Request, RequestStatus
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
|
||||
|
||||
# isort: on
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_update_sched(scheduler_cls) -> MagicMock:
|
||||
"""MagicMock scheduler driving the real ``update_from_output`` with an
|
||||
empty batch; only the fence state carries real values."""
|
||||
sched = MagicMock()
|
||||
sched.requests = {}
|
||||
sched.perf_metrics = None
|
||||
sched.chunk_transfer_adapter = None
|
||||
sched.defer_block_free = True
|
||||
sched.processed_step_seq = 0
|
||||
return sched
|
||||
|
||||
|
||||
def _empty_frame(*, total_num_scheduled_tokens: int):
|
||||
scheduler_output = MagicMock(spec=SchedulerOutput)
|
||||
scheduler_output.num_scheduled_tokens = {}
|
||||
scheduler_output.total_num_scheduled_tokens = total_num_scheduled_tokens
|
||||
model_runner_output = MagicMock(spec=ModelRunnerOutput)
|
||||
model_runner_output.sampled_token_ids = []
|
||||
model_runner_output.logprobs = None
|
||||
model_runner_output.prompt_logprobs_dict = {}
|
||||
model_runner_output.pooler_output = None
|
||||
model_runner_output.num_nans_in_logits = None
|
||||
model_runner_output.kv_connector_output = None
|
||||
model_runner_output.cudagraph_stats = None
|
||||
return scheduler_output, model_runner_output
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
|
||||
def test_update_from_output_advances_fence_and_drains(scheduler_cls) -> None:
|
||||
sched = _make_update_sched(scheduler_cls)
|
||||
scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=3)
|
||||
|
||||
scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
|
||||
|
||||
assert sched.processed_step_seq == 1
|
||||
sched._drain_deferred_frees.assert_called_once_with()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
|
||||
def test_update_from_output_zero_token_step_leaves_fence(scheduler_cls) -> None:
|
||||
"""0-token steps do not advance the fence upstream (their seq was never
|
||||
advanced in schedule()); draining on them would run ahead of the GPU."""
|
||||
sched = _make_update_sched(scheduler_cls)
|
||||
scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=0)
|
||||
|
||||
scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
|
||||
|
||||
assert sched.processed_step_seq == 0
|
||||
sched._drain_deferred_frees.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
|
||||
def test_update_from_output_defer_disabled_leaves_fence(scheduler_cls) -> None:
|
||||
sched = _make_update_sched(scheduler_cls)
|
||||
sched.defer_block_free = False
|
||||
scheduler_output, model_runner_output = _empty_frame(total_num_scheduled_tokens=3)
|
||||
|
||||
scheduler_cls.update_from_output(sched, scheduler_output, model_runner_output)
|
||||
|
||||
assert sched.processed_step_seq == 0
|
||||
sched._drain_deferred_frees.assert_not_called()
|
||||
|
||||
|
||||
def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request]:
|
||||
"""MagicMock scheduler that drives the real generation fast-path
|
||||
``schedule()`` over one running one-shot request."""
|
||||
request = Request(
|
||||
request_id="req-fence",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
sampling_params=SamplingParams(max_tokens=4),
|
||||
pooling_params=None,
|
||||
arrival_time=100.0,
|
||||
block_hasher=None,
|
||||
)
|
||||
request.status = RequestStatus.RUNNING
|
||||
|
||||
sched = MagicMock()
|
||||
sched.max_num_scheduled_tokens = 8
|
||||
sched._pause_state = PauseState.UNPAUSED
|
||||
sched.requests = {request.request_id: request}
|
||||
sched.running = [request]
|
||||
sched.waiting = [] # falsy: the waiting loop never runs
|
||||
sched.policy = SchedulingPolicy.FCFS
|
||||
sched.chunk_transfer_adapter = None
|
||||
sched.scheduler_config = MagicMock(enable_chunked_prefill=True)
|
||||
sched.num_lookahead_tokens = 0
|
||||
sched.log_stats = False
|
||||
sched.kv_cache_config = MagicMock(kv_cache_groups=[])
|
||||
sched.use_v2_model_runner = False
|
||||
sched.prev_step_scheduled_req_ids = set()
|
||||
sched.needs_kv_cache_zeroing = False
|
||||
sched.connector = None
|
||||
sched.ec_connector = None
|
||||
sched.defer_block_free = defer_block_free
|
||||
sched.sched_step_seq = 0
|
||||
# last_sched_seq stamping order: _update_after_schedule must observe the
|
||||
# already-advanced fence value, as upstream schedule() guarantees.
|
||||
sched.seq_at_update_after_schedule = None
|
||||
|
||||
def _capture(_scheduler_output):
|
||||
sched.seq_at_update_after_schedule = sched.sched_step_seq
|
||||
|
||||
sched._update_after_schedule.side_effect = _capture
|
||||
return sched, request
|
||||
|
||||
|
||||
def test_fast_path_schedule_advances_fence_before_update_after_schedule() -> None:
|
||||
sched, request = _make_fast_path_sched(defer_block_free=True)
|
||||
|
||||
OmniGenerationScheduler.schedule(sched)
|
||||
|
||||
assert sched.sched_step_seq == 1
|
||||
assert sched.seq_at_update_after_schedule == 1
|
||||
|
||||
|
||||
def test_fast_path_schedule_defer_disabled_leaves_fence() -> None:
|
||||
sched, request = _make_fast_path_sched(defer_block_free=False)
|
||||
|
||||
OmniGenerationScheduler.schedule(sched)
|
||||
|
||||
assert sched.sched_step_seq == 0
|
||||
|
||||
|
||||
def test_generation_schedule_sweeps_and_resyncs_around_pending_inputs() -> None:
|
||||
sched, request = _make_fast_path_sched(defer_block_free=False)
|
||||
calls: list[str] = []
|
||||
sched._drop_aborted_queued_requests.side_effect = lambda: calls.append("sweep")
|
||||
sched._process_pending_omni_inputs.side_effect = lambda model_mode: calls.append(model_mode)
|
||||
sched._resync_streaming_input_counter.side_effect = lambda: calls.append("resync")
|
||||
|
||||
OmniGenerationScheduler.schedule(sched)
|
||||
|
||||
assert calls == ["sweep", "generation", "sweep", "resync"]
|
||||
@@ -0,0 +1,180 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""v0.28 EC request-finish contract in the Omni schedulers (#6606 review).
|
||||
|
||||
Upstream v0.28 ``_free_request()`` fires ``ec_connector.request_finished()``
|
||||
BEFORE the encoder cache is freed (so the connector can inspect per-request
|
||||
state), honors its delayed-free result, and returns the generated
|
||||
``ec_transfer_params``; ``update_from_output()`` then carries them into the
|
||||
``EngineCoreOutput``. ``OmniARScheduler`` previously skipped the hook and
|
||||
always returned ``None``; both Omni schedulers discarded the second tuple
|
||||
item, so a remote encoder-cache handle could never reach the frontend.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
|
||||
# before Request / RequestStatus are bound in this module.
|
||||
# isort: off
|
||||
import vllm_omni # noqa: F401 - import for side effects (patch vLLM)
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.outputs import ModelRunnerOutput
|
||||
from vllm.v1.request import Request, RequestStatus
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
|
||||
|
||||
# isort: on
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
EC_PARAMS = {"remote_encoder_cache_handle": "ec-1"}
|
||||
|
||||
|
||||
def _make_free_request_sched(*, ec_connector) -> tuple[OmniARScheduler, MagicMock]:
|
||||
"""Minimal real-`_free_request` harness; `order` records the EC hook
|
||||
firing relative to the encoder-cache free."""
|
||||
order = MagicMock()
|
||||
sched = OmniARScheduler.__new__(OmniARScheduler)
|
||||
sched._omits_kv_transfer_cache = {}
|
||||
sched._connector_finished = lambda request: (False, None)
|
||||
sched.ec_connector = ec_connector
|
||||
if ec_connector is not None:
|
||||
order.attach_mock(ec_connector.request_finished, "request_finished")
|
||||
sched.encoder_cache_manager = MagicMock()
|
||||
order.attach_mock(sched.encoder_cache_manager.free, "encoder_cache_free")
|
||||
sched.finished_req_ids = set()
|
||||
sched._new_prompt_len_snapshot = {}
|
||||
sched.finished_req_ids_dict = None
|
||||
sched._should_transfer_kv_for_request = lambda req_id: False
|
||||
sched._free_blocks = MagicMock()
|
||||
sched._free_input_coordinator_request = MagicMock()
|
||||
sched.chunk_transfer_adapter = None
|
||||
return sched, order
|
||||
|
||||
|
||||
class _FakeFinishedRequest:
|
||||
def __init__(self, request_id: str) -> None:
|
||||
self.request_id = request_id
|
||||
|
||||
def is_finished(self) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
def test_free_request_fires_ec_hook_before_encoder_cache_free() -> None:
|
||||
ec_connector = MagicMock()
|
||||
ec_connector.request_finished.return_value = (False, EC_PARAMS)
|
||||
sched, order = _make_free_request_sched(ec_connector=ec_connector)
|
||||
request = _FakeFinishedRequest("req-ec")
|
||||
|
||||
kv_params, ec_params = sched._free_request(request)
|
||||
|
||||
assert ec_params == EC_PARAMS
|
||||
ec_connector.request_finished.assert_called_once_with(request)
|
||||
hook_calls = [name for name, *_ in order.mock_calls]
|
||||
assert hook_calls.index("request_finished") < hook_calls.index("encoder_cache_free")
|
||||
sched._free_blocks.assert_called_once_with(request)
|
||||
|
||||
|
||||
def test_free_request_honors_ec_delayed_free() -> None:
|
||||
ec_connector = MagicMock()
|
||||
ec_connector.request_finished.return_value = (True, EC_PARAMS)
|
||||
sched, _ = _make_free_request_sched(ec_connector=ec_connector)
|
||||
request = _FakeFinishedRequest("req-ec-delay")
|
||||
|
||||
kv_params, ec_params = sched._free_request(request)
|
||||
|
||||
assert ec_params == EC_PARAMS
|
||||
sched._free_blocks.assert_not_called()
|
||||
|
||||
|
||||
def test_free_request_without_ec_connector_returns_none() -> None:
|
||||
sched, _ = _make_free_request_sched(ec_connector=None)
|
||||
request = _FakeFinishedRequest("req-no-ec")
|
||||
|
||||
kv_params, ec_params = sched._free_request(request)
|
||||
|
||||
assert ec_params is None
|
||||
sched._free_blocks.assert_called_once_with(request)
|
||||
|
||||
|
||||
def _make_finishing_request() -> Request:
|
||||
request = Request(
|
||||
request_id="req-ec-finish",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
sampling_params=SamplingParams(max_tokens=4),
|
||||
pooling_params=None,
|
||||
arrival_time=100.0,
|
||||
block_hasher=None,
|
||||
)
|
||||
request.status = RequestStatus.RUNNING
|
||||
return request
|
||||
|
||||
|
||||
def _make_finish_sched(request: Request) -> MagicMock:
|
||||
"""MagicMock scheduler driving the real ``update_from_output`` through a
|
||||
finishing request; ``_free_request`` reports EC transfer params."""
|
||||
sched = MagicMock()
|
||||
sched.requests = {request.request_id: request}
|
||||
sched.perf_metrics = None
|
||||
sched.defer_block_free = False
|
||||
sched.structured_output_manager.should_advance.return_value = False
|
||||
sched._process_kv_transfer_trigger.return_value = False
|
||||
sched._maybe_decode_pooling_output.return_value = None
|
||||
sched._handle_stopped_request.return_value = True
|
||||
sched._free_request.return_value = (None, EC_PARAMS)
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
sched.chunk_transfer_adapter = None
|
||||
sched.finished_req_ids_dict = None
|
||||
sched._new_prompt_len_snapshot = {}
|
||||
return sched
|
||||
|
||||
|
||||
def test_ar_update_from_output_carries_ec_params_to_engine_output() -> None:
|
||||
request = _make_finishing_request()
|
||||
sched = _make_finish_sched(request)
|
||||
sched._update_request_with_output.return_value = ([42], True)
|
||||
|
||||
engine_core_outputs = OmniARScheduler.update_from_output(sched, *_finish_frame_parts(request))
|
||||
|
||||
(eco,) = engine_core_outputs[request.client_index].outputs
|
||||
assert eco.ec_transfer_params == EC_PARAMS
|
||||
sched._free_request.assert_called_once_with(request)
|
||||
|
||||
|
||||
def _finish_frame_parts(request: Request):
|
||||
scheduler_output = MagicMock(spec=SchedulerOutput)
|
||||
scheduler_output.num_scheduled_tokens = {request.request_id: 1}
|
||||
scheduler_output.total_num_scheduled_tokens = 1
|
||||
scheduler_output.scheduled_spec_decode_tokens = {}
|
||||
scheduler_output.num_invalid_spec_tokens = 0
|
||||
|
||||
model_runner_output = MagicMock(spec=ModelRunnerOutput)
|
||||
model_runner_output.sampled_token_ids = [[42]]
|
||||
model_runner_output.logprobs = None
|
||||
model_runner_output.prompt_logprobs_dict = {}
|
||||
model_runner_output.pooler_output = None
|
||||
model_runner_output.num_nans_in_logits = None
|
||||
model_runner_output.kv_connector_output = None
|
||||
model_runner_output.cudagraph_stats = None
|
||||
model_runner_output.req_id_to_index = {request.request_id: 0}
|
||||
model_runner_output.routed_experts = None
|
||||
return scheduler_output, model_runner_output
|
||||
|
||||
|
||||
def test_generation_update_from_output_carries_ec_params_to_engine_output() -> None:
|
||||
request = _make_finishing_request()
|
||||
# One-shot generation finish: the whole prompt is computed.
|
||||
request.num_computed_tokens = len(request.prompt_token_ids)
|
||||
sched = _make_finish_sched(request)
|
||||
|
||||
engine_core_outputs = OmniGenerationScheduler.update_from_output(sched, *_finish_frame_parts(request))
|
||||
|
||||
(eco,) = engine_core_outputs[request.client_index].outputs
|
||||
assert eco.ec_transfer_params == EC_PARAMS
|
||||
sched._free_request.assert_called_once_with(request)
|
||||
@@ -0,0 +1,126 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""v0.28 prefill-stats finalization in the Omni schedulers (#6606 review).
|
||||
|
||||
Upstream v0.28 finalizes prefill statistics with
|
||||
``PrefillStats.finalize(kv_cache_manager.estimate_cached_tokens(request))``
|
||||
before emitting them, which computes ``num_cache_creation_tokens`` (tokens
|
||||
computed and written to the prefix cache). Both Omni schedulers previously
|
||||
emitted ``take_prefill_stats()`` raw, so creation usage was always 0.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
# Imports must run in this order: vllm_omni applies patches to vllm.v1.request
|
||||
# before Request / RequestStatus are bound in this module.
|
||||
# isort: off
|
||||
import vllm_omni # noqa: F401 - import for side effects (patch vLLM)
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.outputs import ModelRunnerOutput
|
||||
from vllm.v1.request import Request, RequestStatus
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
|
||||
|
||||
# isort: on
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_request() -> Request:
|
||||
request = Request(
|
||||
request_id="req-prefill-stats",
|
||||
prompt_token_ids=[1, 2, 3],
|
||||
sampling_params=SamplingParams(max_tokens=4),
|
||||
pooling_params=None,
|
||||
arrival_time=100.0,
|
||||
block_hasher=None,
|
||||
)
|
||||
request.status = RequestStatus.RUNNING
|
||||
# A scheduled prefill: 3 prompt tokens, 1 from the local prefix cache.
|
||||
request.prefill_stats.set(
|
||||
num_prompt_tokens=3,
|
||||
num_local_cached_tokens=1,
|
||||
num_external_cached_tokens=0,
|
||||
)
|
||||
return request
|
||||
|
||||
|
||||
def _make_sched(request: Request) -> MagicMock:
|
||||
sched = MagicMock()
|
||||
sched.requests = {request.request_id: request}
|
||||
sched.perf_metrics = None
|
||||
sched.defer_block_free = False
|
||||
sched.structured_output_manager.should_advance.return_value = False
|
||||
sched._process_kv_transfer_trigger.return_value = False
|
||||
sched._maybe_decode_pooling_output.return_value = None
|
||||
sched._handle_stopped_request.return_value = True
|
||||
# The whole 3-token prompt ends up cached: 3 - 1 pre-cached = 2 created.
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 3
|
||||
|
||||
def _free_and_invalidate(req):
|
||||
# Lifecycle guard: freeing releases the KV blocks, after which the
|
||||
# manager reports 0 cached tokens. If finalize ran after
|
||||
# _free_request (the reviewed bug), creation count would compute
|
||||
# against 0 and the assertions below would see 0, not 2.
|
||||
sched.kv_cache_manager.estimate_cached_tokens.return_value = 0
|
||||
return (None, None)
|
||||
|
||||
sched._free_request.side_effect = _free_and_invalidate
|
||||
sched.chunk_transfer_adapter = None
|
||||
sched.finished_req_ids_dict = None
|
||||
sched._new_prompt_len_snapshot = {}
|
||||
return sched
|
||||
|
||||
|
||||
def _frame(request: Request):
|
||||
scheduler_output = MagicMock(spec=SchedulerOutput)
|
||||
scheduler_output.num_scheduled_tokens = {request.request_id: 1}
|
||||
scheduler_output.total_num_scheduled_tokens = 1
|
||||
scheduler_output.scheduled_spec_decode_tokens = {}
|
||||
scheduler_output.num_invalid_spec_tokens = 0
|
||||
|
||||
model_runner_output = MagicMock(spec=ModelRunnerOutput)
|
||||
model_runner_output.sampled_token_ids = [[42]]
|
||||
model_runner_output.logprobs = None
|
||||
model_runner_output.prompt_logprobs_dict = {}
|
||||
model_runner_output.pooler_output = None
|
||||
model_runner_output.num_nans_in_logits = None
|
||||
model_runner_output.kv_connector_output = None
|
||||
model_runner_output.cudagraph_stats = None
|
||||
model_runner_output.req_id_to_index = {request.request_id: 0}
|
||||
model_runner_output.routed_experts = None
|
||||
return scheduler_output, model_runner_output
|
||||
|
||||
|
||||
def test_ar_scheduler_finalizes_prefill_stats_before_emitting() -> None:
|
||||
request = _make_request()
|
||||
sched = _make_sched(request)
|
||||
sched._update_request_with_output.return_value = ([42], True)
|
||||
|
||||
engine_core_outputs = OmniARScheduler.update_from_output(sched, *_frame(request))
|
||||
|
||||
(eco,) = engine_core_outputs[request.client_index].outputs
|
||||
assert eco.prefill_stats is not None
|
||||
assert eco.prefill_stats.num_cached_tokens == 1
|
||||
assert eco.prefill_stats.num_cache_creation_tokens == 2
|
||||
sched.kv_cache_manager.estimate_cached_tokens.assert_called_once_with(request)
|
||||
|
||||
|
||||
def test_generation_scheduler_finalizes_prefill_stats_before_emitting() -> None:
|
||||
request = _make_request()
|
||||
request.num_computed_tokens = len(request.prompt_token_ids)
|
||||
sched = _make_sched(request)
|
||||
|
||||
engine_core_outputs = OmniGenerationScheduler.update_from_output(sched, *_frame(request))
|
||||
|
||||
(eco,) = engine_core_outputs[request.client_index].outputs
|
||||
assert eco.prefill_stats is not None
|
||||
assert eco.prefill_stats.num_cached_tokens == 1
|
||||
assert eco.prefill_stats.num_cache_creation_tokens == 2
|
||||
sched.kv_cache_manager.estimate_cached_tokens.assert_called_once_with(request)
|
||||
@@ -157,3 +157,71 @@ def test_finish_requests_leaves_healthy_running_intact(
|
||||
)
|
||||
|
||||
assert scheduler.running == [alive]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
|
||||
def test_finish_requests_preserves_unrelated_resumable_segment(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
scheduler_cls,
|
||||
scheduler_mod,
|
||||
) -> None:
|
||||
paused = _StubRequest("req-paused", RequestStatus.FINISHED_STOPPED)
|
||||
paused.resumable = True
|
||||
leaving = _StubRequest("req-leaving", RequestStatus.FINISHED_STOPPED)
|
||||
scheduler = _make_scheduler(
|
||||
scheduler_cls,
|
||||
requests={paused.request_id: paused},
|
||||
running=[paused],
|
||||
waiting=[],
|
||||
)
|
||||
|
||||
def fake_finish_requests(self, request_ids, finished_status):
|
||||
return [leaving]
|
||||
|
||||
monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
|
||||
|
||||
scheduler_cls.finish_requests(scheduler, [leaving.request_id], RequestStatus.FINISHED_ABORTED)
|
||||
|
||||
assert scheduler.running == [paused]
|
||||
assert scheduler.requests == {paused.request_id: paused}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
|
||||
def test_finish_requests_materializes_single_pass_request_ids(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
scheduler_cls,
|
||||
scheduler_mod,
|
||||
) -> None:
|
||||
request = _StubRequest("req-generator", RequestStatus.RUNNING)
|
||||
scheduler = _make_scheduler(
|
||||
scheduler_cls,
|
||||
requests={request.request_id: request},
|
||||
running=[request],
|
||||
waiting=[],
|
||||
)
|
||||
|
||||
def fake_finish_requests(self, request_ids, finished_status):
|
||||
assert tuple(request_ids) == (request.request_id,)
|
||||
return [request]
|
||||
|
||||
monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
|
||||
|
||||
result = scheduler_cls.finish_requests(
|
||||
scheduler,
|
||||
(request_id for request_id in [request.request_id]),
|
||||
RequestStatus.FINISHED_ABORTED,
|
||||
)
|
||||
|
||||
assert result == [request]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
|
||||
def test_finish_requests_does_not_reopen_off_queue_deferred_free_terminal(scheduler_cls) -> None:
|
||||
terminal = _StubRequest("req-deferred-free", RequestStatus.FINISHED_STOPPED)
|
||||
terminal.resumable = True
|
||||
scheduler = _make_scheduler(scheduler_cls, requests={terminal.request_id: terminal}, running=[], waiting=[])
|
||||
scheduler._free_request = lambda *args, **kwargs: pytest.fail("off-queue request was freed")
|
||||
|
||||
assert scheduler_cls.finish_requests(scheduler, [terminal.request_id], RequestStatus.FINISHED_ABORTED) == []
|
||||
assert scheduler.requests == {terminal.request_id: terminal}
|
||||
assert terminal.status == RequestStatus.FINISHED_STOPPED
|
||||
@@ -4,6 +4,7 @@ from types import SimpleNamespace
|
||||
import pytest
|
||||
from vllm.v1.engine import FinishReason
|
||||
|
||||
from vllm_omni.core.sched import omni_scheduler_mixin
|
||||
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
|
||||
from vllm_omni.core.sched.output import OmniChunkRecvHandle
|
||||
|
||||
@@ -14,6 +15,77 @@ class _Scheduler(OmniSchedulerMixin):
|
||||
pass
|
||||
|
||||
|
||||
def test_async_chunk_adapter_initializes_for_stage_zero_sender_and_stage_one_receiver(monkeypatch):
|
||||
created_adapters = []
|
||||
|
||||
def adapter_factory(config):
|
||||
adapter = SimpleNamespace(vllm_config=config)
|
||||
created_adapters.append(adapter)
|
||||
return adapter
|
||||
|
||||
monkeypatch.setattr(omni_scheduler_mixin, "OmniChunkTransferAdapter", adapter_factory)
|
||||
|
||||
def init_scheduler(**model_config):
|
||||
scheduler = _Scheduler()
|
||||
scheduler.vllm_config = SimpleNamespace(model_config=SimpleNamespace(**model_config))
|
||||
scheduler._init_omni_io_scheduling_state()
|
||||
return scheduler
|
||||
|
||||
producer = init_scheduler(
|
||||
stage_id=0,
|
||||
async_chunk=True,
|
||||
requires_full_payload_input=False,
|
||||
custom_process_next_stage_input_func="test.pipeline.produce_async_chunk",
|
||||
)
|
||||
receiver = init_scheduler(
|
||||
stage_id=1,
|
||||
async_chunk=True,
|
||||
requires_full_payload_input=True,
|
||||
custom_process_next_stage_input_func=None,
|
||||
)
|
||||
terminal_stage = init_scheduler(
|
||||
stage_id=2,
|
||||
async_chunk=True,
|
||||
requires_full_payload_input=False,
|
||||
custom_process_next_stage_input_func=None,
|
||||
)
|
||||
|
||||
assert producer.chunk_transfer_adapter is not None
|
||||
assert receiver.chunk_transfer_adapter is not None
|
||||
assert terminal_stage.chunk_transfer_adapter is not None
|
||||
assert receiver.input_coordinator is None
|
||||
assert len(created_adapters) == 3
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("stage_id", "async_chunk", "required", "enabled"),
|
||||
[
|
||||
(0, False, True, False),
|
||||
(1, True, True, False),
|
||||
(1, False, False, False),
|
||||
(1, False, True, True),
|
||||
],
|
||||
)
|
||||
def test_full_payload_coordinator_matches_legacy_gate(monkeypatch, stage_id, async_chunk, required, enabled):
|
||||
monkeypatch.setattr(
|
||||
omni_scheduler_mixin,
|
||||
"OmniChunkTransferAdapter",
|
||||
lambda _config: SimpleNamespace(),
|
||||
)
|
||||
scheduler = _Scheduler()
|
||||
scheduler.vllm_config = SimpleNamespace(
|
||||
model_config=SimpleNamespace(
|
||||
stage_id=stage_id,
|
||||
async_chunk=async_chunk,
|
||||
requires_full_payload_input=required,
|
||||
)
|
||||
)
|
||||
|
||||
scheduler._init_omni_io_scheduling_state()
|
||||
|
||||
assert (scheduler.input_coordinator is not None) is enabled
|
||||
|
||||
|
||||
def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
|
||||
calls = []
|
||||
scheduler = _Scheduler()
|
||||
@@ -22,13 +94,25 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
|
||||
scheduler.requests = {"request": object()}
|
||||
scheduler._consume_pending_connector_output = lambda mode: calls.append(("consume", mode))
|
||||
scheduler._process_pending_input_timeouts = lambda: calls.append(("timeouts",))
|
||||
|
||||
def _collect_timed_out(timeout_s):
|
||||
calls.append(("chunk-timeouts", timeout_s))
|
||||
return set()
|
||||
|
||||
def _collect_failed_sends():
|
||||
calls.append(("failed-sends",))
|
||||
return {}
|
||||
|
||||
scheduler.chunk_transfer_adapter = SimpleNamespace(
|
||||
receives_chunks=True,
|
||||
process_pending_chunks=lambda waiting, running, scheduler_requests: calls.append(
|
||||
("process", waiting, running, scheduler_requests)
|
||||
),
|
||||
restore_queues=lambda waiting, running, scheduler_requests: calls.append(
|
||||
("restore-chunks", waiting, running, scheduler_requests)
|
||||
),
|
||||
collect_timed_out_request_ids=_collect_timed_out,
|
||||
collect_failed_send_request_ids=_collect_failed_sends,
|
||||
)
|
||||
scheduler.input_coordinator = SimpleNamespace(
|
||||
restore_queues=lambda waiting: calls.append(("restore-full", waiting))
|
||||
@@ -41,6 +125,10 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
|
||||
("consume", "ar"),
|
||||
("timeouts",),
|
||||
("process", scheduler.waiting, scheduler.running, scheduler.requests),
|
||||
# The chunk deadline runs after chunks are applied, so a chunk that
|
||||
# arrived this cycle resets the clock before it is measured (R1.1).
|
||||
("chunk-timeouts", omni_scheduler_mixin.DEFAULT_INPUT_WAIT_TIMEOUT_S),
|
||||
("failed-sends",),
|
||||
("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests),
|
||||
("restore-full", scheduler.waiting),
|
||||
]
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit coverage for _process_pending_input_timeouts.
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Unit coverage for the two stage-input timeout hooks.
|
||||
|
||||
``_process_pending_input_timeouts`` covers the full-payload transport;
|
||||
``_process_pending_chunk_timeouts`` covers async-chunk. Both share
|
||||
``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S``.
|
||||
|
||||
Verifies that the mixin correctly *delegates* timed-out requests to the
|
||||
base scheduler's ``finish_requests`` API with ``RequestStatus.FINISHED_ERROR``.
|
||||
@@ -12,6 +16,10 @@ this file only asserts the wiring from the mixin to that API.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
@@ -102,5 +110,280 @@ def test_process_pending_input_timeouts_disabled_when_timeout_zero(monkeypatch):
|
||||
assert scheduler.finish_calls == []
|
||||
|
||||
|
||||
class _FakeChunkAdapter:
|
||||
"""Stands in for OmniChunkTransferAdapter's timeout-collection surface."""
|
||||
|
||||
def __init__(self, timed_out_ids, *, receives_chunks=True):
|
||||
self._timed_out_ids = set(timed_out_ids)
|
||||
self.receives_chunks = receives_chunks
|
||||
self.calls = []
|
||||
|
||||
def collect_timed_out_request_ids(self, timeout_s):
|
||||
self.calls.append(timeout_s)
|
||||
return set(self._timed_out_ids)
|
||||
|
||||
|
||||
class _FakeChunkScheduler(OmniSchedulerMixin):
|
||||
def __init__(self, requests, adapter):
|
||||
self.requests = requests
|
||||
self.chunk_transfer_adapter = adapter
|
||||
self.finish_calls = []
|
||||
|
||||
def finish_requests(self, req_ids, status):
|
||||
self.finish_calls.append((set(req_ids), status))
|
||||
|
||||
|
||||
def test_process_pending_chunk_timeouts_delegates_to_finish_requests():
|
||||
"""A request stalled in WAITING_FOR_CHUNK is failed, not left parked.
|
||||
|
||||
This is the async-chunk half of the safety net (RFC #4855 R1.1); before it
|
||||
existed, the chunk path had no deadline at all and issue #3833 could park a
|
||||
request forever.
|
||||
"""
|
||||
req_id = "stalled-stream"
|
||||
adapter = _FakeChunkAdapter(timed_out_ids={req_id})
|
||||
scheduler = _FakeChunkScheduler({req_id: SimpleNamespace(request_id=req_id)}, adapter)
|
||||
|
||||
scheduler._process_pending_chunk_timeouts()
|
||||
|
||||
assert len(adapter.calls) == 1
|
||||
assert adapter.calls[0] > 0
|
||||
assert len(scheduler.finish_calls) == 1
|
||||
finished_ids, status = scheduler.finish_calls[0]
|
||||
assert finished_ids == {req_id}
|
||||
assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR")
|
||||
|
||||
|
||||
def test_process_pending_chunk_timeouts_skips_already_freed_request():
|
||||
adapter = _FakeChunkAdapter(timed_out_ids={"already-freed"})
|
||||
scheduler = _FakeChunkScheduler(requests={}, adapter=adapter)
|
||||
|
||||
scheduler._process_pending_chunk_timeouts()
|
||||
|
||||
assert scheduler.finish_calls == []
|
||||
|
||||
|
||||
def test_process_pending_chunk_timeouts_noop_on_a_producer_stage():
|
||||
"""Stage 0 does not receive chunks, so it has nothing to expire."""
|
||||
adapter = _FakeChunkAdapter(timed_out_ids={"r1"}, receives_chunks=False)
|
||||
scheduler = _FakeChunkScheduler({"r1": SimpleNamespace(request_id="r1")}, adapter)
|
||||
|
||||
scheduler._process_pending_chunk_timeouts()
|
||||
|
||||
assert adapter.calls == []
|
||||
assert scheduler.finish_calls == []
|
||||
|
||||
|
||||
def test_process_pending_chunk_timeouts_noop_without_adapter():
|
||||
class _NoAdapter(OmniSchedulerMixin):
|
||||
def __init__(self):
|
||||
self.requests = {}
|
||||
self.chunk_transfer_adapter = None
|
||||
self.finish_calls = []
|
||||
|
||||
def finish_requests(self, req_ids, status):
|
||||
self.finish_calls.append((set(req_ids), status))
|
||||
|
||||
scheduler = _NoAdapter()
|
||||
scheduler._process_pending_chunk_timeouts()
|
||||
assert scheduler.finish_calls == []
|
||||
|
||||
|
||||
def test_process_pending_chunk_timeouts_disabled_when_timeout_zero(monkeypatch):
|
||||
"""The shared knob disables both transports' nets together."""
|
||||
from vllm_omni.core.sched import omni_scheduler_mixin
|
||||
|
||||
monkeypatch.setattr(omni_scheduler_mixin, "DEFAULT_INPUT_WAIT_TIMEOUT_S", 0.0)
|
||||
|
||||
adapter = _FakeChunkAdapter(timed_out_ids={"r1"})
|
||||
scheduler = _FakeChunkScheduler({"r1": SimpleNamespace(request_id="r1")}, adapter)
|
||||
scheduler._process_pending_chunk_timeouts()
|
||||
|
||||
assert adapter.calls == []
|
||||
assert scheduler.finish_calls == []
|
||||
|
||||
|
||||
class _FakeSendFailureAdapter:
|
||||
"""Stands in for the adapter's drained send-failure map (R1.2)."""
|
||||
|
||||
def __init__(self, failures):
|
||||
self._failures = dict(failures)
|
||||
self.collect_calls = 0
|
||||
|
||||
def collect_failed_send_request_ids(self):
|
||||
self.collect_calls += 1
|
||||
failures, self._failures = self._failures, {}
|
||||
return failures
|
||||
|
||||
|
||||
class _FakeSendFailureScheduler(OmniSchedulerMixin):
|
||||
def __init__(self, requests, adapter):
|
||||
self.requests = requests
|
||||
self.chunk_transfer_adapter = adapter
|
||||
|
||||
|
||||
def test_log_failed_chunk_sends_reports_a_live_request(caplog):
|
||||
adapter = _FakeSendFailureAdapter({"req-live": "OSError: No space left on device"})
|
||||
scheduler = _FakeSendFailureScheduler({"req-live": SimpleNamespace(request_id="req-live")}, adapter)
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
scheduler._log_failed_chunk_sends()
|
||||
|
||||
assert "req-live" in caplog.text
|
||||
assert "No space left on device" in caplog.text
|
||||
|
||||
|
||||
def test_log_failed_chunk_sends_still_reports_a_departed_request(caplog):
|
||||
"""``collect_*`` drains the map, so a failure for a request that already left
|
||||
this scheduler is the one place a drop could still go unrecorded."""
|
||||
adapter = _FakeSendFailureAdapter({"req-gone": "OSError: No space left on device"})
|
||||
scheduler = _FakeSendFailureScheduler({}, adapter)
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
scheduler._log_failed_chunk_sends()
|
||||
|
||||
assert adapter.collect_calls == 1
|
||||
assert "req-gone" in caplog.text
|
||||
assert "No space left on device" in caplog.text
|
||||
assert "finished or aborted" in caplog.text
|
||||
|
||||
|
||||
def test_log_failed_chunk_sends_noop_without_adapter():
|
||||
scheduler = _FakeSendFailureScheduler({}, None)
|
||||
scheduler._log_failed_chunk_sends()
|
||||
|
||||
|
||||
class _FakeReceiveFailureAdapter:
|
||||
def __init__(self, failures):
|
||||
self._failures = dict(failures)
|
||||
|
||||
def collect_failed_receive_request_ids(self):
|
||||
failures, self._failures = self._failures, {}
|
||||
return failures
|
||||
|
||||
|
||||
class _FakeReceiveFailureScheduler(OmniSchedulerMixin):
|
||||
def __init__(self, requests, adapter):
|
||||
self.requests = requests
|
||||
self.chunk_transfer_adapter = adapter
|
||||
self.finish_calls = []
|
||||
|
||||
def finish_requests(self, req_ids, status):
|
||||
self.finish_calls.append((set(req_ids), status))
|
||||
|
||||
|
||||
def test_process_chunk_receive_failures_delegates_to_finish_requests():
|
||||
request_id = "invalid-chunk"
|
||||
adapter = _FakeReceiveFailureAdapter({request_id: "prompt exceeds limit"})
|
||||
scheduler = _FakeReceiveFailureScheduler({request_id: SimpleNamespace(request_id=request_id)}, adapter)
|
||||
|
||||
scheduler._process_chunk_receive_failures()
|
||||
|
||||
assert len(scheduler.finish_calls) == 1
|
||||
finished_ids, status = scheduler.finish_calls[0]
|
||||
assert finished_ids == {request_id}
|
||||
assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR")
|
||||
|
||||
|
||||
# --- R1.4: VLLM_OMNI_INPUT_WAIT_TIMEOUT_S is parsed at import time, so each of
|
||||
# these reloads the module under a different environment. One knob now gates the
|
||||
# deadline on both transports, which is why a silent 0 matters more than it did
|
||||
# when only the full-payload path used it.
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def reload_with(monkeypatch):
|
||||
"""Reload the mixin module under a given env value, then put it back.
|
||||
|
||||
``importlib.reload`` re-executes into the live module dict, so whatever the
|
||||
last test parsed is what every later test in the session sees -- ``-k
|
||||
zero_warns`` alone would leave ``DEFAULT_INPUT_WAIT_TIMEOUT_S = 0.0`` and
|
||||
silently disable both deadlines. The rejection cases are worse: the
|
||||
ValueError fires *after* the rebind, so ``nan`` lingers in the module.
|
||||
``monkeypatch`` restores the environment but not the module, so reload once
|
||||
more on the way out.
|
||||
"""
|
||||
import vllm_omni.core.sched.omni_scheduler_mixin as mod
|
||||
|
||||
def _reload(raw):
|
||||
if raw is None:
|
||||
monkeypatch.delenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", raising=False)
|
||||
else:
|
||||
monkeypatch.setenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", raw)
|
||||
return importlib.reload(mod)
|
||||
|
||||
yield _reload
|
||||
|
||||
monkeypatch.undo()
|
||||
importlib.reload(mod)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"raw,expected",
|
||||
[
|
||||
(None, 600.0),
|
||||
("30", 30.0),
|
||||
("0.5", 0.5),
|
||||
("0", 0.0), # explicit opt-out is still honoured
|
||||
],
|
||||
)
|
||||
def test_timeout_accepts_valid_values(reload_with, raw, expected):
|
||||
mod = reload_with(raw)
|
||||
assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raw", ["-1", "-600", "-0.5"])
|
||||
def test_negative_timeout_is_rejected(reload_with, raw):
|
||||
"""`-1` is a common "no limit" idiom; substituting 600 would give an operator
|
||||
who meant "never time out" mysterious failures ten minutes in. Fail loudly."""
|
||||
with pytest.raises(ValueError, match="negative"):
|
||||
reload_with(raw)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("raw", ["nan", "inf", "-inf", "1e999"])
|
||||
def test_non_finite_timeout_is_rejected(reload_with, raw):
|
||||
"""nan compares false against every deadline check and inf is unreachable;
|
||||
both disable the net while looking like a number."""
|
||||
with pytest.raises(ValueError, match="not finite"):
|
||||
reload_with(raw)
|
||||
|
||||
|
||||
def test_zero_warns_and_names_both_transports(reload_with, caplog):
|
||||
with caplog.at_level(logging.WARNING):
|
||||
mod = reload_with("0")
|
||||
assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == 0.0
|
||||
text = caplog.text
|
||||
assert "DISABLED" in text
|
||||
assert "WAITING_FOR_INPUT" in text and "WAITING_FOR_CHUNK" in text
|
||||
|
||||
|
||||
def test_unparsable_timeout_falls_back(reload_with):
|
||||
mod = reload_with("not-a-number")
|
||||
assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == 600.0
|
||||
|
||||
|
||||
def test_a_rejected_value_does_not_outlive_its_test(monkeypatch):
|
||||
"""The reason ``reload_with`` reloads on the way out, done here by hand.
|
||||
|
||||
``nan`` is bound before the ValueError fires, so the reload that rejects it
|
||||
still leaves it in the module -- and nan compares false against every
|
||||
deadline check, which is exactly the silent-disable this item exists to
|
||||
stop. Restoring the environment is not enough; the module has to be reloaded
|
||||
too.
|
||||
"""
|
||||
import vllm_omni.core.sched.omni_scheduler_mixin as mod
|
||||
|
||||
ambient = float(os.environ.get("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "600"))
|
||||
|
||||
monkeypatch.setenv("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "nan")
|
||||
with pytest.raises(ValueError, match="not finite"):
|
||||
importlib.reload(mod)
|
||||
assert math.isnan(mod.DEFAULT_INPUT_WAIT_TIMEOUT_S)
|
||||
|
||||
monkeypatch.undo()
|
||||
importlib.reload(mod)
|
||||
assert mod.DEFAULT_INPUT_WAIT_TIMEOUT_S == ambient
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -0,0 +1,168 @@
|
||||
"""Tests for the ``num_waiting_for_streaming_input`` resync placed in
|
||||
``OmniARScheduler.finish_requests`` / ``OmniGenerationScheduler.finish_requests``
|
||||
and in ``OmniARScheduler.schedule``.
|
||||
|
||||
Upstream keeps that counter incrementally (``+1`` when a resumable request
|
||||
parks in a waiting queue as ``WAITING_FOR_STREAMING_REQ``, ``-1`` when the
|
||||
next update arrives or the request is finished) and
|
||||
``Scheduler.get_num_unfinished_requests`` *subtracts* it from the waiting
|
||||
queues. Omni rewrites ``request.status`` outside both hooks -- the
|
||||
chunk-transfer adapter's park/restore and
|
||||
``_realign_request_status_to_queues`` -- so a counted request can leave
|
||||
``WAITING_FOR_STREAMING_REQ`` without the matching decrement and inflate the
|
||||
counter for the lifetime of the process.
|
||||
|
||||
An inflated counter does not degrade throughput, it hangs the stage:
|
||||
``EngineCore.has_work()`` reads false while a live request sits in
|
||||
``waiting``, so the engine blocks in ``input_queue.get()`` and never calls
|
||||
``schedule()`` again. Regression for the Qwen3-Omni duplex symptom where the
|
||||
server answered the first conversation after a boot and was silent from the
|
||||
second on while still reporting healthy.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
|
||||
from vllm.v1.request import RequestStatus
|
||||
|
||||
import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
|
||||
import vllm_omni.core.sched.omni_generation_scheduler as gen_sched_mod
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _StubRequest:
|
||||
def __init__(self, request_id: str, status: RequestStatus) -> None:
|
||||
self.request_id = request_id
|
||||
self.status = status
|
||||
|
||||
def is_finished(self) -> bool:
|
||||
return RequestStatus.is_finished(self.status)
|
||||
|
||||
|
||||
_SCHEDULER_PARAMS = [
|
||||
pytest.param(OmniARScheduler, ar_sched_mod, id="ar"),
|
||||
pytest.param(OmniGenerationScheduler, gen_sched_mod, id="generation"),
|
||||
]
|
||||
|
||||
|
||||
def _make_scheduler(scheduler_cls, *, requests, running, waiting, counter, skipped=None):
|
||||
scheduler = scheduler_cls.__new__(scheduler_cls)
|
||||
scheduler.chunk_transfer_adapter = None
|
||||
scheduler.input_coordinator = None
|
||||
scheduler.requests = requests
|
||||
scheduler.running = running
|
||||
scheduler.waiting = waiting
|
||||
scheduler.skipped_waiting = skipped if skipped is not None else []
|
||||
scheduler.num_waiting_for_streaming_input = counter
|
||||
return scheduler
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
|
||||
def test_finish_requests_clears_counter_left_by_a_stomped_status(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
scheduler_cls,
|
||||
scheduler_mod,
|
||||
) -> None:
|
||||
"""The abort case that wedged the talker.
|
||||
|
||||
The request was counted while parked for streaming input, then had its
|
||||
status rewritten to ``RUNNING`` by omni's queue surgery, so upstream's
|
||||
``finish_requests`` took the running branch and never decremented. With
|
||||
no parked requests left, the counter must land at 0.
|
||||
"""
|
||||
aborted = _StubRequest("req-aborted", RequestStatus.RUNNING)
|
||||
scheduler = _make_scheduler(
|
||||
scheduler_cls,
|
||||
requests={}, # upstream already dropped it
|
||||
running=[],
|
||||
waiting=[],
|
||||
counter=1,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
scheduler_mod.VLLMScheduler,
|
||||
"finish_requests",
|
||||
lambda self, request_ids, finished_status: [aborted],
|
||||
)
|
||||
|
||||
scheduler_cls.finish_requests(scheduler, [aborted.request_id], RequestStatus.FINISHED_ABORTED)
|
||||
|
||||
assert scheduler.num_waiting_for_streaming_input == 0
|
||||
|
||||
|
||||
def test_resync_counts_parked_requests_in_skipped_waiting() -> None:
|
||||
"""``get_num_unfinished_requests`` sums ``waiting`` and ``skipped_waiting``,
|
||||
so the resync has to look at both or it under-counts and reintroduces the
|
||||
spin the counter exists to prevent."""
|
||||
parked = _StubRequest("req-skipped", RequestStatus.WAITING_FOR_STREAMING_REQ)
|
||||
scheduler = _make_scheduler(
|
||||
OmniARScheduler,
|
||||
requests={parked.request_id: parked},
|
||||
running=[],
|
||||
waiting=[],
|
||||
counter=0,
|
||||
skipped=[parked],
|
||||
)
|
||||
|
||||
scheduler._resync_streaming_input_counter()
|
||||
|
||||
assert scheduler.num_waiting_for_streaming_input == 1
|
||||
|
||||
|
||||
def test_ar_schedule_resyncs_before_delegating_upstream(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Pins placement in ``schedule``: the chunk-transfer adapter rewrites
|
||||
status during ``process_pending_chunks``, so the resync has to run after
|
||||
that and before upstream schedules, while ``self.waiting`` is still the
|
||||
real queue."""
|
||||
parked = _StubRequest("req-resumed", RequestStatus.RUNNING)
|
||||
scheduler = OmniARScheduler.__new__(OmniARScheduler)
|
||||
scheduler.chunk_transfer_adapter = None
|
||||
scheduler.input_coordinator = None
|
||||
scheduler.requests = {parked.request_id: parked}
|
||||
scheduler.running = [parked]
|
||||
scheduler.waiting = []
|
||||
scheduler.skipped_waiting = []
|
||||
scheduler.num_waiting_for_streaming_input = 1 # leaked by the status rewrite
|
||||
scheduler.policy = "fcfs"
|
||||
|
||||
seen: dict[str, int] = {}
|
||||
|
||||
monkeypatch.setattr(
|
||||
ar_sched_mod.OmniSchedulerMixin, "_consume_pending_connector_output", lambda self, model_mode: None
|
||||
)
|
||||
monkeypatch.setattr(ar_sched_mod.OmniSchedulerMixin, "_process_pending_input_timeouts", lambda self: None)
|
||||
monkeypatch.setattr(OmniARScheduler, "_should_defer_waiting_admission", lambda self: False)
|
||||
|
||||
def fake_schedule(self, throttle_prefills=False):
|
||||
seen["counter"] = self.num_waiting_for_streaming_input
|
||||
raise _StopScheduleError
|
||||
|
||||
class _StopScheduleError(Exception):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", fake_schedule)
|
||||
|
||||
with pytest.raises(_StopScheduleError):
|
||||
OmniARScheduler.schedule(scheduler)
|
||||
|
||||
assert seen["counter"] == 0
|
||||
|
||||
|
||||
def test_generation_scheduler_sweep_removes_aborted_requests_from_all_queues() -> None:
|
||||
aborted = _StubRequest("req-generation-aborted", RequestStatus.FINISHED_ABORTED)
|
||||
scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
|
||||
scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
|
||||
scheduler.waiting.add_request(aborted)
|
||||
scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS)
|
||||
scheduler.skipped_waiting.add_request(aborted)
|
||||
scheduler.running = [aborted]
|
||||
|
||||
scheduler._drop_aborted_queued_requests()
|
||||
|
||||
assert list(scheduler.waiting) == []
|
||||
assert list(scheduler.skipped_waiting) == []
|
||||
assert scheduler.running == []
|
||||
@@ -20,7 +20,6 @@ import torch
|
||||
import vllm_omni.core.sched.omni_scheduling_coordinator as coord_mod
|
||||
from vllm_omni.core.sched.omni_scheduling_coordinator import (
|
||||
OmniSchedulingCoordinator,
|
||||
uses_full_payload_input_coordinator,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -96,109 +95,6 @@ class MockQueue:
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
|
||||
class TestFullPayloadCoordinatorSelection(unittest.TestCase):
|
||||
"""Tests for the (model_arch, model_stage) whitelist gate.
|
||||
|
||||
The init_omni_connectors arch allowlist is keyed by ``model_arch`` and
|
||||
is a superset of the stages registered here -- consumer-wait stages
|
||||
must be registered explicitly in ``_FULL_PAYLOAD_INPUT_STAGES``, while
|
||||
the init allowlist covers both producer- and consumer-side runners.
|
||||
These tests pin which ``(arch, stage)`` pairs the gate fires for today.
|
||||
"""
|
||||
|
||||
# Expected whitelist (model_arch, model_stage). Hardcoded to avoid the
|
||||
# tautology of importing _FULL_PAYLOAD_INPUT_STAGES and asserting it
|
||||
# against itself; any drift between this matrix and the whitelist will
|
||||
# fail loudly here.
|
||||
EXPECTED_FULL_PAYLOAD_INPUT_STAGES: frozenset[tuple[str, str]] = frozenset(
|
||||
{
|
||||
("Qwen3OmniMoeForConditionalGeneration", "talker"),
|
||||
("Qwen3OmniMoeForConditionalGeneration", "code2wav"),
|
||||
("Qwen2_5OmniForConditionalGeneration", "talker"),
|
||||
("Qwen2_5OmniForConditionalGeneration", "code2wav"),
|
||||
("CovoAudioForConditionalGeneration", "code2wav"),
|
||||
("MiMoAudioModel", "code2wav"),
|
||||
("Qwen3TTSCode2Wav", "code2wav"),
|
||||
("MiniCPMO45Code2Wav", "code2wav"),
|
||||
# nemotron_voicechat: sync-path RVQ-VAE decode waits for the
|
||||
# talker's full [frames, 31] code-stack payload.
|
||||
("NemotronVoiceChatCode2Wav", "code2wav"),
|
||||
("CosyVoice3Model", "cosyvoice3_code2wav"),
|
||||
# audex: sync-path TTS decoder (async_chunk: false deploys) and
|
||||
# the TTA XCodec1 stage (always sync full-payload).
|
||||
("AudexCode2Wav", "audex_code2wav"),
|
||||
("AudexXCodec1", "audex_xcodec"),
|
||||
("IndexTTS2S2MelDecoder", "indextts2_s2mel_decoder"),
|
||||
("IndexTTS25S2MelDecoder", "indextts2_5_s2mel_decoder"),
|
||||
("DyninOmniForConditionalGeneration", "token2image"),
|
||||
("DyninOmniForConditionalGeneration", "token2audio"),
|
||||
}
|
||||
)
|
||||
|
||||
def test_whitelist_matches_expected_matrix(self):
|
||||
"""_FULL_PAYLOAD_INPUT_STAGES must equal the hardcoded expected matrix.
|
||||
|
||||
Catches both accidental additions (which would silently enable the
|
||||
consumer-wait gate for a new arch) and accidental removals (which
|
||||
would silently disable an enabled arch).
|
||||
"""
|
||||
from vllm_omni.core.sched.omni_scheduling_coordinator import _FULL_PAYLOAD_INPUT_STAGES
|
||||
|
||||
self.assertEqual(
|
||||
frozenset(_FULL_PAYLOAD_INPUT_STAGES),
|
||||
self.EXPECTED_FULL_PAYLOAD_INPUT_STAGES,
|
||||
msg="_FULL_PAYLOAD_INPUT_STAGES drifted from the expected matrix; "
|
||||
"update EXPECTED_FULL_PAYLOAD_INPUT_STAGES if intentional.",
|
||||
)
|
||||
|
||||
def test_all_whitelisted_arch_stage_pairs_fire_gate(self):
|
||||
"""Every (arch, stage) pair in the expected matrix must fire
|
||||
the gate when stage_id > 0 and async_chunk=False.
|
||||
"""
|
||||
for arch, stage in self.EXPECTED_FULL_PAYLOAD_INPUT_STAGES:
|
||||
model_config = SimpleNamespace(
|
||||
stage_id=1,
|
||||
async_chunk=False,
|
||||
model_arch=arch,
|
||||
model_stage=stage,
|
||||
)
|
||||
self.assertTrue(
|
||||
uses_full_payload_input_coordinator(model_config),
|
||||
msg=f"expected gate to fire for {arch}/{stage}",
|
||||
)
|
||||
|
||||
def test_other_arch_or_stage_or_mode_does_not_fire(self):
|
||||
cases = [
|
||||
SimpleNamespace(
|
||||
stage_id=1, async_chunk=True, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage="talker"
|
||||
),
|
||||
SimpleNamespace(
|
||||
stage_id=0, async_chunk=False, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage="thinker"
|
||||
),
|
||||
SimpleNamespace(
|
||||
stage_id=1,
|
||||
async_chunk=False,
|
||||
model_arch="Qwen3OmniMoeForConditionalGeneration",
|
||||
model_stage="some_future_stage",
|
||||
),
|
||||
SimpleNamespace(
|
||||
stage_id=1, async_chunk=False, model_arch="Qwen3TTSForConditionalGeneration", model_stage="code2wav"
|
||||
),
|
||||
SimpleNamespace(
|
||||
stage_id=1, async_chunk=False, model_arch="MingFlashOmniForConditionalGeneration", model_stage="talker"
|
||||
),
|
||||
SimpleNamespace(stage_id=1, async_chunk=False, model_arch=None, model_stage="talker"),
|
||||
SimpleNamespace(
|
||||
stage_id=1, async_chunk=False, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage=None
|
||||
),
|
||||
]
|
||||
for model_config in cases:
|
||||
self.assertFalse(
|
||||
uses_full_payload_input_coordinator(model_config),
|
||||
msg=f"expected gate OFF for {model_config}",
|
||||
)
|
||||
|
||||
|
||||
class TestCoordinatorUpdateRequestMetadata(unittest.TestCase):
|
||||
"""Test update_request_metadata applies scheduling metadata to requests."""
|
||||
|
||||
|
||||
@@ -105,7 +105,7 @@ def create_unique_server_pytest_params(
|
||||
configs: list[dict[str, Any]],
|
||||
stage_configs_dir: Path,
|
||||
) -> list[Any]:
|
||||
"""Like :func:`create_unique_server_params`, but wrap each row in ``pytest.param`` with JSON marks."""
|
||||
"""Like :func:`_create_unique_server_params`, but wrap each row in ``pytest.param`` with JSON marks."""
|
||||
marks_by_name = _marks_by_test_name(configs)
|
||||
return [
|
||||
pytest.param(
|
||||
@@ -113,7 +113,7 @@ def create_unique_server_pytest_params(
|
||||
marks=marks_by_name.get(row[0], []),
|
||||
id=row[0],
|
||||
)
|
||||
for row in create_unique_server_params(configs, stage_configs_dir)
|
||||
for row in _create_unique_server_params(configs, stage_configs_dir)
|
||||
]
|
||||
|
||||
|
||||
@@ -150,7 +150,7 @@ def create_paired_omni_benchmark_pytest_params(
|
||||
"""Paired params for ``run_benchmark.py`` (omni/tts)."""
|
||||
mapping = create_test_parameter_mapping(configs)
|
||||
marks_by_name = _marks_by_test_name(configs)
|
||||
server_entries = [(row, row[0]) for row in create_unique_server_params(configs, stage_configs_dir)]
|
||||
server_entries = [(row, row[0]) for row in _create_unique_server_params(configs, stage_configs_dir)]
|
||||
params_by_test_name = {
|
||||
test_name: get_benchmark_params_for_server(test_name, mapping) for _, test_name in server_entries
|
||||
}
|
||||
@@ -232,7 +232,7 @@ def _build_serve_args(serve_args: Any) -> list[str]:
|
||||
return args
|
||||
|
||||
|
||||
def create_unique_server_params(
|
||||
def _create_unique_server_params(
|
||||
configs: list[dict[str, Any]],
|
||||
stage_configs_dir: Path,
|
||||
) -> list[tuple[str, str, str | None, str | None, tuple[str, ...], bool]]:
|
||||
@@ -311,7 +311,7 @@ def create_reliability_omni_server_params(
|
||||
configs: list[dict[str, Any]], stage_configs_dir: Path
|
||||
) -> list[OmniServerParams]:
|
||||
adjusted_configs = configs_with_platform_stage_configs(configs)
|
||||
unique_params = create_unique_server_params(adjusted_configs, stage_configs_dir)
|
||||
unique_params = _create_unique_server_params(adjusted_configs, stage_configs_dir)
|
||||
server_args_by_name = extract_server_args_by_test_name(adjusted_configs)
|
||||
return [
|
||||
OmniServerParams(
|
||||
@@ -653,6 +653,7 @@ def run_benchmark(
|
||||
random_input_len: Any | None = None,
|
||||
random_output_len: Any | None = None,
|
||||
resource_label: str | None = None,
|
||||
num_warmups: int = 2,
|
||||
) -> dict[str, Any]:
|
||||
"""Run one ``vllm bench serve --omni`` iteration and return parsed metrics.
|
||||
|
||||
@@ -677,7 +678,7 @@ def run_benchmark(
|
||||
+ args
|
||||
+ [
|
||||
"--num-warmups",
|
||||
"2",
|
||||
str(num_warmups),
|
||||
"--save-result",
|
||||
"--result-dir",
|
||||
os.environ.get("BENCHMARK_DIR", "tests"),
|
||||
|
||||
@@ -1,4 +1,8 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import threading
|
||||
from collections.abc import Callable
|
||||
@@ -159,6 +163,18 @@ def benchmark_params(request):
|
||||
|
||||
def assert_result(result, params, num_prompt) -> None:
|
||||
assert result["completed"] == num_prompt, "Request failures exist"
|
||||
baseline = params.get("baseline")
|
||||
hardware = result.get("Hardware")
|
||||
hardware_baseline = baseline.get(hardware) if isinstance(baseline, dict) and isinstance(hardware, str) else None
|
||||
if isinstance(hardware_baseline, dict) and "mean_tpot_ms" in hardware_baseline:
|
||||
num_tpot_samples = result.get("num_tpot_samples")
|
||||
mean_tpot_ms = result.get("mean_tpot_ms")
|
||||
assert isinstance(num_tpot_samples, int) and not isinstance(num_tpot_samples, bool) and num_tpot_samples > 0, (
|
||||
"TPOT baseline is configured, but no measurable TPOT samples were produced"
|
||||
)
|
||||
assert (
|
||||
isinstance(mean_tpot_ms, int | float) and not isinstance(mean_tpot_ms, bool) and math.isfinite(mean_tpot_ms)
|
||||
), "TPOT baseline is configured, but mean_tpot_ms is not finite"
|
||||
expected_audio_turns = params.get("expected_duplex_audio_turns_per_session")
|
||||
if expected_audio_turns is not None:
|
||||
session_metrics = result.get("duplex_session_metrics")
|
||||
@@ -279,5 +295,6 @@ def test_performance_benchmark(omni_server, benchmark_params):
|
||||
random_input_len=params.get("random_input_len"),
|
||||
random_output_len=params.get("random_output_len"),
|
||||
resource_label=resource_label,
|
||||
num_warmups=max(2, int(concurrency)),
|
||||
)
|
||||
assert_result(result, params, num_prompt)
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"local_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-stage BAGEL (TP=1, CFG=1), t2i Task",
|
||||
@@ -43,8 +43,8 @@
|
||||
"H100": {
|
||||
"throughput_qps": 0.495,
|
||||
"latency_mean": 2.0238,
|
||||
"peak_memory_mb_max": 29393.4286,
|
||||
"peak_memory_mb_mean": 29392.7143
|
||||
"peak_memory_mb_max": 29392.0,
|
||||
"peak_memory_mb_mean": 29392.0
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -61,7 +61,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"local_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-stage BAGEL (TP=1, CFG=1), i2i Task",
|
||||
@@ -93,10 +93,10 @@
|
||||
"num-input-images": 1,
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"throughput_qps": 0.3451,
|
||||
"latency_mean": 2.9031,
|
||||
"peak_memory_mb_max": 30480.2857,
|
||||
"peak_memory_mb_mean": 30479.2857
|
||||
"throughput_qps": 0.3499,
|
||||
"latency_mean": 2.8701,
|
||||
"peak_memory_mb_max": 30466.0,
|
||||
"peak_memory_mb_mean": 30465.0
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -113,7 +113,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"local_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, t2i Task",
|
||||
@@ -164,7 +164,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"local_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, i2i Task",
|
||||
@@ -196,8 +196,8 @@
|
||||
"num-input-images": 1,
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"throughput_qps": 0.149,
|
||||
"latency_mean": 6.7157,
|
||||
"throughput_qps": 0.1533,
|
||||
"latency_mean": 6.5243,
|
||||
"peak_memory_mb_max": 30574.0,
|
||||
"peak_memory_mb_mean": 30574.0
|
||||
}
|
||||
|
||||
@@ -49,8 +49,8 @@
|
||||
},
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"throughput_qps": 1.107,
|
||||
"latency_mean": 0.8929,
|
||||
"throughput_qps": 1.1958,
|
||||
"latency_mean": 0.843,
|
||||
"peak_memory_mb_mean": 80000.0
|
||||
}
|
||||
}
|
||||
@@ -118,9 +118,9 @@
|
||||
},
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"throughput_qps": 0.027,
|
||||
"latency_mean": 36.6824,
|
||||
"peak_memory_mb_mean": 30155.4
|
||||
"throughput_qps": 0.0307,
|
||||
"latency_mean": 32.7665,
|
||||
"peak_memory_mb_mean": 28762.4762
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -186,9 +186,9 @@
|
||||
},
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"throughput_qps": 0.027,
|
||||
"latency_mean": 36.968,
|
||||
"peak_memory_mb_mean": 30443.6
|
||||
"throughput_qps": 0.0305,
|
||||
"latency_mean": 33.1009,
|
||||
"peak_memory_mb_mean": 29441.1429
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -258,9 +258,9 @@
|
||||
},
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"throughput_qps": 0.027,
|
||||
"latency_mean": 37.2117,
|
||||
"peak_memory_mb_mean": 30135.6
|
||||
"throughput_qps": 0.0299,
|
||||
"latency_mean": 33.925,
|
||||
"peak_memory_mb_mean": 29443.1429
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"local_model",
|
||||
"tts"
|
||||
],
|
||||
"server_params": {
|
||||
|
||||
@@ -137,9 +137,7 @@
|
||||
"H200": {
|
||||
"throughput_qps": 0.1874,
|
||||
"latency_p99": 101.852,
|
||||
"latency_mean": 79.0385,
|
||||
"peak_memory_mb_max": 70000,
|
||||
"peak_memory_mb_mean": 70000
|
||||
"latency_mean": 79.0385
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,106 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_hunyuanvideo15_i2v_single_device",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H100"
|
||||
},
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device baseline",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v",
|
||||
"serve_args": {
|
||||
"flow-shift": 5.0,
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "832x480_frames33_steps4",
|
||||
"dataset": "random",
|
||||
"task": "i2v",
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"num-input-images": 1,
|
||||
"seed": 42,
|
||||
"enable-negative-prompt": true,
|
||||
"random-request-config": [
|
||||
{
|
||||
"width": 832,
|
||||
"height": 480,
|
||||
"num_inference_steps": 4,
|
||||
"num_frames": 33,
|
||||
"fps": 24,
|
||||
"weight": 1
|
||||
}
|
||||
],
|
||||
"extra-body": {
|
||||
"guidance_scale": 6.0,
|
||||
"flow_shift": 5.0
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_hunyuanvideo15_i2v_cachedit_tp2_vae_patch2",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H100"
|
||||
},
|
||||
"num_cards": 2
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "CacheDiT + TP=2 + VAE patch parallel=2 + VAE tiling",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v",
|
||||
"serve_args": {
|
||||
"flow-shift": 5.0,
|
||||
"cache-backend": "cache_dit",
|
||||
"tensor-parallel-size": 2,
|
||||
"vae-patch-parallel-size": 2,
|
||||
"vae-use-tiling": true,
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "832x480_frames33_steps4",
|
||||
"dataset": "random",
|
||||
"task": "i2v",
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"num-input-images": 1,
|
||||
"seed": 42,
|
||||
"enable-negative-prompt": true,
|
||||
"random-request-config": [
|
||||
{
|
||||
"width": 832,
|
||||
"height": 480,
|
||||
"num_inference_steps": 4,
|
||||
"num_frames": 33,
|
||||
"fps": 24,
|
||||
"weight": 1
|
||||
}
|
||||
],
|
||||
"extra-body": {
|
||||
"guidance_scale": 6.0,
|
||||
"flow_shift": 5.0
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -0,0 +1,104 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_hunyuanvideo15_t2v_single_device",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H100"
|
||||
},
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device baseline",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
|
||||
"serve_args": {
|
||||
"flow-shift": 5.0,
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "832x480_frames33_steps4",
|
||||
"dataset": "random",
|
||||
"task": "t2v",
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"seed": 42,
|
||||
"enable-negative-prompt": true,
|
||||
"random-request-config": [
|
||||
{
|
||||
"width": 832,
|
||||
"height": 480,
|
||||
"num_inference_steps": 4,
|
||||
"num_frames": 33,
|
||||
"fps": 24,
|
||||
"weight": 1
|
||||
}
|
||||
],
|
||||
"extra-body": {
|
||||
"guidance_scale": 6.0,
|
||||
"flow_shift": 5.0
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_hunyuanvideo15_t2v_cachedit_tp2_vae_patch2",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H100"
|
||||
},
|
||||
"num_cards": 2
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "CacheDiT + TP=2 + VAE patch parallel=2 + VAE tiling",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v",
|
||||
"serve_args": {
|
||||
"flow-shift": 5.0,
|
||||
"cache-backend": "cache_dit",
|
||||
"tensor-parallel-size": 2,
|
||||
"vae-patch-parallel-size": 2,
|
||||
"vae-use-tiling": true,
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "832x480_frames33_steps4",
|
||||
"dataset": "random",
|
||||
"task": "t2v",
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"seed": 42,
|
||||
"enable-negative-prompt": true,
|
||||
"random-request-config": [
|
||||
{
|
||||
"width": 832,
|
||||
"height": 480,
|
||||
"num_inference_steps": 4,
|
||||
"num_frames": 33,
|
||||
"fps": 24,
|
||||
"weight": 1
|
||||
}
|
||||
],
|
||||
"extra-body": {
|
||||
"guidance_scale": 6.0,
|
||||
"flow_shift": 5.0
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"local_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device baseline with enforce-eager (no torch.compile)",
|
||||
@@ -64,7 +64,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"local_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device with torch.compile (default, no enforce-eager)",
|
||||
@@ -117,7 +117,7 @@
|
||||
"num_cards": 2
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"local_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "CFG-parallel=2 with enforce-eager",
|
||||
@@ -172,7 +172,7 @@
|
||||
"num_cards": 2
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"local_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "CFG-parallel=2 with torch.compile",
|
||||
@@ -226,7 +226,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"local_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "CacheDiT with enforce-eager",
|
||||
|
||||
@@ -39,6 +39,7 @@
|
||||
8
|
||||
],
|
||||
"no_oversample": true,
|
||||
"disable_shuffle": true,
|
||||
"trust_remote_code": true,
|
||||
"seed_tts_locale": "en",
|
||||
"extra_body": {
|
||||
@@ -55,33 +56,116 @@
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"request_throughput": [
|
||||
0.7803,
|
||||
1.1008,
|
||||
1.2027
|
||||
0.9731,
|
||||
1.1986,
|
||||
1.686
|
||||
],
|
||||
"mean_ttft_ms": [
|
||||
166.9764,
|
||||
307.7133,
|
||||
323.584
|
||||
152.5665,
|
||||
270.2559,
|
||||
301.455
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
1361.1682,
|
||||
3800.6332,
|
||||
6961.8636
|
||||
1027.2953,
|
||||
3306.3735,
|
||||
4686.5366
|
||||
],
|
||||
"mean_audio_ttfp_ms": [
|
||||
635.8099,
|
||||
1615.6591,
|
||||
2888.9831
|
||||
458.1553,
|
||||
2030.2335,
|
||||
2068.6233
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.3151,
|
||||
0.8568,
|
||||
1.5425
|
||||
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
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_minicpmo_4_5_a2",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"npu": "A2"
|
||||
},
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"omni"
|
||||
],
|
||||
"server_params": {
|
||||
"model": "openbmb/MiniCPM-o-4_5",
|
||||
"extra_cli_args": [
|
||||
"--deploy-config",
|
||||
"vllm_omni/deploy/minicpmo_4_5.yaml",
|
||||
"--trust-remote-code"
|
||||
]
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"dataset_name": "seed-tts",
|
||||
"dataset_path": "zhaochenyang20/seed-tts-eval",
|
||||
"backend": "openai-chat-omni",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"num_prompts": [
|
||||
32,
|
||||
64,
|
||||
128
|
||||
],
|
||||
"max_concurrency": [
|
||||
1,
|
||||
4,
|
||||
8
|
||||
],
|
||||
"no_oversample": true,
|
||||
"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"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -18,7 +18,7 @@
|
||||
"model": "openbmb/MiniCPM-o-4_5",
|
||||
"extra_cli_args": [
|
||||
"--deploy-config",
|
||||
"vllm_omni/deploy/minicpmo_4_5_duplex.yaml",
|
||||
"vllm_omni/deploy/minicpmo_4_5.yaml",
|
||||
"--trust-remote-code"
|
||||
]
|
||||
},
|
||||
@@ -35,6 +35,7 @@
|
||||
1
|
||||
],
|
||||
"no_oversample": true,
|
||||
"disable_shuffle": true,
|
||||
"trust_remote_code": true,
|
||||
"seed_tts_locale": "en",
|
||||
"seed_tts_turns_per_session": 4,
|
||||
@@ -42,9 +43,26 @@
|
||||
"extra_body": {
|
||||
"save_duplex_request_metrics": true
|
||||
},
|
||||
"percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration",
|
||||
"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
|
||||
],
|
||||
|
||||
@@ -37,6 +37,7 @@
|
||||
4
|
||||
],
|
||||
"no_oversample": true,
|
||||
"disable_shuffle": true,
|
||||
"trust_remote_code": true,
|
||||
"seed_tts_locale": "en",
|
||||
"extra_body": {
|
||||
@@ -53,24 +54,24 @@
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"request_throughput": [
|
||||
0.7803,
|
||||
1.1008
|
||||
0.9731,
|
||||
1.1986
|
||||
],
|
||||
"mean_ttft_ms": [
|
||||
166.9764,
|
||||
307.7133
|
||||
152.5665,
|
||||
270.2559
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
1361.1682,
|
||||
3800.6332
|
||||
1027.2953,
|
||||
3306.3735
|
||||
],
|
||||
"mean_audio_ttfp_ms": [
|
||||
635.8099,
|
||||
1615.6591
|
||||
458.1553,
|
||||
2030.2335
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.3151,
|
||||
0.8568
|
||||
0.2357,
|
||||
0.791
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
@@ -21,54 +21,6 @@
|
||||
]
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"dataset_name": "random",
|
||||
"backend": "openai-chat-omni",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"num_prompts": [
|
||||
4,
|
||||
16,
|
||||
32,
|
||||
64,
|
||||
128
|
||||
],
|
||||
"max_concurrency": [
|
||||
1,
|
||||
4,
|
||||
8,
|
||||
16,
|
||||
32
|
||||
],
|
||||
"random_input_len": 2500,
|
||||
"random_output_len": 900,
|
||||
"ignore_eos": true,
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
96.4061,
|
||||
140.8147,
|
||||
271.9025,
|
||||
362.2916,
|
||||
507.776
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
18507.4582,
|
||||
28365.4956,
|
||||
31907.9613,
|
||||
48161.5829,
|
||||
72630.8601
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1175,
|
||||
0.18,
|
||||
0.2359,
|
||||
0.3278,
|
||||
0.6296
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"dataset_name": "random-mm",
|
||||
"backend": "openai-chat-omni",
|
||||
@@ -94,14 +46,8 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
354.0951
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
3952.5615
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.117
|
||||
0.1098
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -133,14 +79,8 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
259.6317
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
4253.3009
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1448
|
||||
0.1335
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -174,60 +114,8 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
457.4453
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
5912.5623
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1967
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"dataset_name": "random",
|
||||
"backend": "openai-chat-omni",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"num_prompts": [
|
||||
4,
|
||||
16,
|
||||
32,
|
||||
64,
|
||||
128
|
||||
],
|
||||
"max_concurrency": [
|
||||
1,
|
||||
4,
|
||||
8,
|
||||
16,
|
||||
32
|
||||
],
|
||||
"random_input_len": 2500,
|
||||
"random_output_len": 900,
|
||||
"ignore_eos": true,
|
||||
"extra_body": {
|
||||
"modalities": [
|
||||
"text"
|
||||
]
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
90.8033,
|
||||
202.1779,
|
||||
326.5864,
|
||||
544.8914,
|
||||
961.3929
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
4517.3411,
|
||||
7857.9711,
|
||||
9861.1687,
|
||||
13283.3363,
|
||||
19541.0736
|
||||
0.1798
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -262,11 +150,8 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
318.048
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
807.6114
|
||||
756.8415
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -303,11 +188,8 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
186.168
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
726.1137
|
||||
694.9777
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -346,9 +228,6 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
375.5282
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
1006.1663
|
||||
]
|
||||
@@ -356,5 +235,61 @@
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_qwen3_omni_chunk_cuda",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H100"
|
||||
},
|
||||
"num_cards": 2
|
||||
}
|
||||
},
|
||||
"omni",
|
||||
"slow"
|
||||
],
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
||||
"extra_cli_args": [
|
||||
"--async-chunk"
|
||||
]
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"dataset_name": "random",
|
||||
"backend": "openai-chat-omni",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"num_prompts": [
|
||||
4,
|
||||
32,
|
||||
128
|
||||
],
|
||||
"max_concurrency": [
|
||||
1,
|
||||
8,
|
||||
32
|
||||
],
|
||||
"random_input_len": 100,
|
||||
"random_output_len": 100,
|
||||
"ignore_eos": true,
|
||||
"extra_body": {
|
||||
"modalities": [
|
||||
"text"
|
||||
]
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_e2el_ms": [
|
||||
4354.8945,
|
||||
9861.1687,
|
||||
18716.709
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 3
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"slow",
|
||||
"omni"
|
||||
],
|
||||
"server_params": {
|
||||
@@ -36,10 +36,9 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_ttft_ms": [130.5751, 278.6469, 324.7249, 911.6085],
|
||||
"mean_tpot_ms": [7.9854, 11.1209, 16.1507, 20.7485],
|
||||
"mean_audio_ttfp_ms": [486.486, 954.1791, 1289.6514, 2196.8275],
|
||||
"mean_audio_rtf": [0.1896, 0.3962, 0.4098, 0.7322]
|
||||
"mean_tpot_ms": [7.6955, 10.6457, 12.7363, 15.1608],
|
||||
"mean_audio_ttfp_ms": [325.1624, 602.5192, 699.3891, 1024.088],
|
||||
"mean_audio_rtf": [0.1689, 0.2088, 0.2397, 0.2791]
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -64,9 +63,8 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_ttft_ms": [340.5122],
|
||||
"mean_audio_ttfp_ms": [463.0752],
|
||||
"mean_audio_rtf": [0.1153]
|
||||
"mean_audio_ttfp_ms": [436.4663],
|
||||
"mean_audio_rtf": [0.1069]
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -93,9 +91,8 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_ttft_ms": [236.926],
|
||||
"mean_audio_ttfp_ms": [395.2182],
|
||||
"mean_audio_rtf": [0.1259]
|
||||
"mean_audio_ttfp_ms": [376.4911],
|
||||
"mean_audio_rtf": [0.1145]
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -124,9 +121,8 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_ttft_ms": [470.6126],
|
||||
"mean_audio_ttfp_ms": [665.5474],
|
||||
"mean_audio_rtf": [0.166]
|
||||
"mean_audio_ttfp_ms": [638.2227],
|
||||
"mean_audio_rtf": [0.1438]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -45,26 +45,12 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
115.8348,
|
||||
173.1823,
|
||||
446.672,
|
||||
464.5985,
|
||||
3420.0634
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
25070.9523,
|
||||
31992.0646,
|
||||
37619.7669,
|
||||
50393.8257,
|
||||
136619.2743
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.169,
|
||||
0.2111,
|
||||
0.2475,
|
||||
0.3551,
|
||||
0.9266
|
||||
0.2041,
|
||||
0.2382,
|
||||
0.3177,
|
||||
0.4595
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -94,14 +80,8 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
343.1377
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
4811.1123
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1463
|
||||
0.138
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -133,14 +113,8 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
245.5446
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
5019.1256
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.1713
|
||||
0.1609
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -174,14 +148,8 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
506.1407
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
6814.9656
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.2213
|
||||
0.2048
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 2
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"slow",
|
||||
"omni"
|
||||
],
|
||||
"server_params": {
|
||||
@@ -29,16 +29,12 @@
|
||||
"endpoint": "/v1/completions",
|
||||
"num_prompts": [
|
||||
4,
|
||||
16,
|
||||
32,
|
||||
64,
|
||||
128
|
||||
],
|
||||
"max_concurrency": [
|
||||
1,
|
||||
4,
|
||||
8,
|
||||
16,
|
||||
32
|
||||
],
|
||||
"random_input_len": 2500,
|
||||
@@ -52,19 +48,10 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
69.8081,
|
||||
173.0651,
|
||||
271.8108,
|
||||
413.4079,
|
||||
497.2546
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
4720.8225,
|
||||
7937.2532,
|
||||
9962.3566,
|
||||
13724.5479,
|
||||
18324.091
|
||||
4619.0068,
|
||||
9760.6175,
|
||||
18007.6918
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -100,53 +87,8 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
25.9477
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
517.3155
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"dataset_name": "random-mm",
|
||||
"backend": "vllm",
|
||||
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
||||
"endpoint": "/v1/completions",
|
||||
"num_prompts": [
|
||||
40
|
||||
],
|
||||
"request_rate": [
|
||||
0.5
|
||||
],
|
||||
"random_input_len": 100,
|
||||
"random_output_len": 100,
|
||||
"random_range_ratio": 0.0,
|
||||
"ignore_eos": true,
|
||||
"extra_body": {
|
||||
"modalities": [
|
||||
"text"
|
||||
]
|
||||
},
|
||||
"random_mm_base_items_per_request": 2,
|
||||
"random_mm_num_mm_items_range_ratio": 0.5,
|
||||
"random_mm_limit_mm_per_prompt": {
|
||||
"image": 1,
|
||||
"video": 1
|
||||
},
|
||||
"random_mm_bucket_config": {
|
||||
"(256, 256, 1)": 0.5,
|
||||
"(720, 1280, 2)": 0.5
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
26.7056
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
551.308
|
||||
491.2534
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -186,11 +128,8 @@
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"mean_ttft_ms": [
|
||||
27.6906
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
596.5284
|
||||
559.6674
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"local_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device baseline (two input images)",
|
||||
@@ -35,10 +35,10 @@
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"throughput_qps": 0.0721,
|
||||
"latency_mean": 13.8398,
|
||||
"peak_memory_mb_max": 57610.0,
|
||||
"peak_memory_mb_mean": 57610.0
|
||||
"throughput_qps": 0.0684,
|
||||
"latency_mean": 14.6249,
|
||||
"peak_memory_mb_max": 57632.0,
|
||||
"peak_memory_mb_mean": 57632.0
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -55,10 +55,10 @@
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"throughput_qps": 0.0186,
|
||||
"latency_mean": 53.6942,
|
||||
"peak_memory_mb_max": 67790.0,
|
||||
"peak_memory_mb_mean": 67790.0
|
||||
"throughput_qps": 0.0177,
|
||||
"latency_mean": 56.5285,
|
||||
"peak_memory_mb_max": 67840.0,
|
||||
"peak_memory_mb_mean": 67840.0
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -76,7 +76,7 @@
|
||||
"num_cards": 4
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"local_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
|
||||
@@ -105,10 +105,10 @@
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"throughput_qps": 0.0538,
|
||||
"latency_mean": 18.5591,
|
||||
"peak_memory_mb_max": 56449.4286,
|
||||
"peak_memory_mb_mean": 56449.4286
|
||||
"throughput_qps": 0.0515,
|
||||
"latency_mean": 19.4345,
|
||||
"peak_memory_mb_max": 56348.0,
|
||||
"peak_memory_mb_mean": 56348.0
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -125,7 +125,7 @@
|
||||
"num_cards": 4
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"local_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Ulysses SP=2 + CFG=2 + CacheDiT",
|
||||
@@ -164,10 +164,10 @@
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"throughput_qps": 0.2776,
|
||||
"latency_mean": 3.5938,
|
||||
"peak_memory_mb_max": 58113.7143,
|
||||
"peak_memory_mb_mean": 58075.1429
|
||||
"throughput_qps": 0.2479,
|
||||
"latency_mean": 4.0389,
|
||||
"peak_memory_mb_max": 57750.0,
|
||||
"peak_memory_mb_mean": 57750.0
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -184,10 +184,10 @@
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"throughput_qps": 0.1051,
|
||||
"latency_mean": 9.4918,
|
||||
"peak_memory_mb_max": 68107.4286,
|
||||
"peak_memory_mb_mean": 68107.4286
|
||||
"throughput_qps": 0.0948,
|
||||
"latency_mean": 10.5462,
|
||||
"peak_memory_mb_max": 67764.0,
|
||||
"peak_memory_mb_mean": 67764.0
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,190 +0,0 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_qwen_image_edit_single_device",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H100"
|
||||
},
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device baseline",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen-Image-Edit",
|
||||
"serve_args": {
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_steps20_i2i",
|
||||
"dataset": "random",
|
||||
"task": "i2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"throughput_qps": 0.1129,
|
||||
"latency_mean": 8.8352,
|
||||
"peak_memory_mb_max": 58572.0,
|
||||
"peak_memory_mb_mean": 58572.0
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "1536x1536_steps35_i2i",
|
||||
"dataset": "random",
|
||||
"task": "i2i",
|
||||
"width": 1536,
|
||||
"height": 1536,
|
||||
"num-inference-steps": 35,
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"throughput_qps": 0.024,
|
||||
"latency_mean": 41.5686,
|
||||
"peak_memory_mb_max": 66368.0,
|
||||
"peak_memory_mb_mean": 66368.0
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_qwen_image_edit_ulysses2_cfg2_vae_patch4",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H100"
|
||||
},
|
||||
"num_cards": 4
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen-Image-Edit",
|
||||
"serve_args": {
|
||||
"ulysses-degree": 2,
|
||||
"cfg-parallel-size": 2,
|
||||
"vae-patch-parallel-size": 4,
|
||||
"vae-use-tiling": true,
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "1536x1536_steps35_i2i",
|
||||
"dataset": "random",
|
||||
"task": "i2i",
|
||||
"width": 1536,
|
||||
"height": 1536,
|
||||
"num-inference-steps": 35,
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"throughput_qps": 0.0672,
|
||||
"latency_mean": 14.8517,
|
||||
"peak_memory_mb_max": 57076.0,
|
||||
"peak_memory_mb_mean": 57076.0
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_qwen_image_edit_ulysses2_cfg2_cache_dit",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H100"
|
||||
},
|
||||
"num_cards": 4
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Ulysses SP=2 + CFG=2 + CacheDiT",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen-Image-Edit",
|
||||
"serve_args": {
|
||||
"ulysses-degree": 2,
|
||||
"cfg-parallel-size": 2,
|
||||
"cache-backend": "cache_dit",
|
||||
"cache-config": {
|
||||
"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": null,
|
||||
"scm_steps_policy": "dynamic"
|
||||
},
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_steps20_i2i",
|
||||
"dataset": "random",
|
||||
"task": "i2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"throughput_qps": 0.2879,
|
||||
"latency_mean": 3.4701,
|
||||
"peak_memory_mb_max": 58758.0,
|
||||
"peak_memory_mb_mean": 58758.0
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "1536x1536_steps35_i2i",
|
||||
"dataset": "random",
|
||||
"task": "i2i",
|
||||
"width": 1536,
|
||||
"height": 1536,
|
||||
"num-inference-steps": 35,
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"throughput_qps": 0.1113,
|
||||
"latency_mean": 8.9656,
|
||||
"peak_memory_mb_max": 67424.0,
|
||||
"peak_memory_mb_mean": 67424.0
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"local_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device baseline",
|
||||
@@ -55,8 +55,8 @@
|
||||
"H100": {
|
||||
"throughput_qps": 0.0417,
|
||||
"latency_mean": 23.914,
|
||||
"peak_memory_mb_max": 63225.1429,
|
||||
"peak_memory_mb_mean": 63225.1429
|
||||
"peak_memory_mb_max": 62912.0,
|
||||
"peak_memory_mb_mean": 62912.0
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -36,7 +36,7 @@
|
||||
"H100": {
|
||||
"throughput_qps": 0.4153,
|
||||
"latency_mean": 2.4065,
|
||||
"peak_memory_mb_mean": 55121.4286
|
||||
"peak_memory_mb_mean": 55062.0
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -54,7 +54,7 @@
|
||||
"H100": {
|
||||
"throughput_qps": 0.0423,
|
||||
"latency_mean": 23.5991,
|
||||
"peak_memory_mb_mean": 65118.0
|
||||
"peak_memory_mb_mean": 64546.0
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -74,13 +74,14 @@
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device baseline (no parallelism) with step execution",
|
||||
"description": "Single-device step execution: sequential 512/1536 plus 512 concurrency sweep. One server (max-num-seqs 8).",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen-Image",
|
||||
"serve_args": {
|
||||
"enable-diffusion-pipeline-profiler": true,
|
||||
"step-execution": true
|
||||
"step-execution": true,
|
||||
"max-num-seqs": 8
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
@@ -98,7 +99,7 @@
|
||||
"H100": {
|
||||
"throughput_qps": 0.4175,
|
||||
"latency_mean": 2.3953,
|
||||
"peak_memory_mb_mean": 55088.5714
|
||||
"peak_memory_mb_mean": 54952.0
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -116,7 +117,28 @@
|
||||
"H100": {
|
||||
"throughput_qps": 0.0423,
|
||||
"latency_mean": 23.5768,
|
||||
"peak_memory_mb_mean": 66012.2857
|
||||
"peak_memory_mb_mean": 64500.0
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "512x512_steps20_high_concurrency",
|
||||
"dataset": "random",
|
||||
"task": "t2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": [20, 40, 80, 160],
|
||||
"max-concurrency": [1, 2, 4, 8],
|
||||
"warmup-requests": 8,
|
||||
"warmup-concurrency": 8,
|
||||
"warmup-num-inference-steps": 20,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"throughput_qps": [0.3937, 0.6606, 0.6188, 0.7752],
|
||||
"latency_mean": [2.5466, 3.0235, 6.3948, 10.1944],
|
||||
"peak_memory_mb_mean": [55149.0, 55150.0, 55153.1458, 55156.325]
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -163,7 +185,7 @@
|
||||
"H100": {
|
||||
"throughput_qps": 0.1214,
|
||||
"latency_mean": 8.2211,
|
||||
"peak_memory_mb_mean": 54501.1429
|
||||
"peak_memory_mb_mean": 54474.0
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -220,7 +242,7 @@
|
||||
"H100": {
|
||||
"throughput_qps": 0.559,
|
||||
"latency_mean": 1.794,
|
||||
"peak_memory_mb_mean": 55290.5714
|
||||
"peak_memory_mb_mean": 55174.0
|
||||
}
|
||||
}
|
||||
},
|
||||
@@ -238,58 +260,10 @@
|
||||
"H100": {
|
||||
"throughput_qps": 0.1918,
|
||||
"latency_mean": 5.2022,
|
||||
"peak_memory_mb_mean": 64838.8571
|
||||
"peak_memory_mb_mean": 64750.0
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_qwen_image_single_device_step_execution_high_concurrency",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H100"
|
||||
},
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"diffusion"
|
||||
],
|
||||
"description": "Single-device Qwen-Image step-level execution benchmark at 512x512, 20 steps, concurrency sweep",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen-Image",
|
||||
"serve_args": {
|
||||
"enable-diffusion-pipeline-profiler": true,
|
||||
"step-execution": true,
|
||||
"max-num-seqs": 8
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_steps20_high_concurrency",
|
||||
"dataset": "random",
|
||||
"task": "t2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": [20, 40, 80, 160],
|
||||
"max-concurrency": [1, 2, 4, 8],
|
||||
"warmup-requests": 8,
|
||||
"warmup-concurrency": 8,
|
||||
"warmup-num-inference-steps": 20,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"throughput_qps": [0.4125, 0.6714, 0.6674, 0.8159],
|
||||
"latency_mean": [2.4215, 2.992, 5.9293, 9.656],
|
||||
"peak_memory_mb_mean": [55147.4286, 55148.0, 55150.4643, 55152.3821]
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -1,3 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Tests for DFX runner metadata field exclusion."""
|
||||
|
||||
import json
|
||||
@@ -105,7 +108,7 @@ def test_resolve_pytest_marks_hardware_dict_with_extra():
|
||||
assert "H100" in names
|
||||
assert "cuda" in names
|
||||
assert "gpu" in names
|
||||
assert "distributed_cuda" in names
|
||||
assert "cards_2" in names
|
||||
assert "full_model" in names
|
||||
assert "diffusion" in names
|
||||
assert "local_model" in names
|
||||
@@ -512,3 +515,41 @@ def test_omni_duplex_expected_audio_turns_rejects_incomplete_session():
|
||||
{"expected_duplex_audio_turns_per_session": 4},
|
||||
1,
|
||||
)
|
||||
|
||||
|
||||
def test_omni_tpot_baseline_accepts_measured_finite_sample():
|
||||
from tests.dfx.perf.scripts.run_benchmark import assert_result
|
||||
|
||||
assert_result(
|
||||
{
|
||||
"completed": 1,
|
||||
"Hardware": "H100",
|
||||
"num_tpot_samples": 1,
|
||||
"mean_tpot_ms": 10.0,
|
||||
},
|
||||
{"baseline": {"H100": {"mean_tpot_ms": 20.0}}},
|
||||
1,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("num_tpot_samples", "mean_tpot_ms", "match"),
|
||||
[
|
||||
(0, float("nan"), "no measurable TPOT samples"),
|
||||
(1, float("nan"), "mean_tpot_ms is not finite"),
|
||||
],
|
||||
)
|
||||
def test_omni_tpot_baseline_rejects_missing_or_nonfinite_sample(num_tpot_samples, mean_tpot_ms, match):
|
||||
from tests.dfx.perf.scripts.run_benchmark import assert_result
|
||||
|
||||
with pytest.raises(AssertionError, match=match):
|
||||
assert_result(
|
||||
{
|
||||
"completed": 1,
|
||||
"Hardware": "H100",
|
||||
"num_tpot_samples": num_tpot_samples,
|
||||
"mean_tpot_ms": mean_tpot_ms,
|
||||
},
|
||||
{"baseline": {"H100": {"mean_tpot_ms": 20.0}}},
|
||||
1,
|
||||
)
|
||||
@@ -424,10 +424,10 @@
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
82.1886
|
||||
82.6344
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1241
|
||||
0.1277
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -453,13 +453,13 @@
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
160.9168
|
||||
170.3737
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1907
|
||||
0.2076
|
||||
],
|
||||
"audio_throughput": [
|
||||
39.2332
|
||||
37.0834
|
||||
]
|
||||
}
|
||||
}
|
||||
@@ -485,13 +485,13 @@
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"median_audio_ttfp_ms": [
|
||||
100.1265
|
||||
101.0293
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.1437
|
||||
0.1498
|
||||
],
|
||||
"audio_throughput": [
|
||||
11.7605
|
||||
10.1186
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"full_model",
|
||||
"local_model",
|
||||
"tts"
|
||||
],
|
||||
"server_params": {
|
||||
|
||||
@@ -7,12 +7,10 @@ The detailed RFC document is maintained outside this repository (local/internal
|
||||
- Qwen3 tests: `test_reliability_qwen3_omni.py`
|
||||
- Wan2.2 tests: `test_reliability_wan22.py`
|
||||
- HunyuanImage tests: `test_reliability_hunyuan_image.py`
|
||||
- VoxCPM2 tests: `test_reliability_voxcpm2.py`
|
||||
|
||||
```bash
|
||||
pytest --collect-only tests/dfx/reliability
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow
|
||||
```
|
||||
@@ -1,13 +1,15 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""``POST /v1/audio/generate`` validation (live Stable Audio Open server)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
|
||||
|
||||
pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
|
||||
|
||||
@@ -95,20 +97,20 @@ _PARAMS = [
|
||||
@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
|
||||
def test_audio_generate_invalid_field_values(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
online_client: OnlineOmniClient,
|
||||
overrides: dict[str, object],
|
||||
err_message: str | tuple[str, ...],
|
||||
) -> None:
|
||||
body = {"model": omni_server.model, "input": "ambient electronic pad"}
|
||||
body: dict[str, Any] = {"model": omni_server.model, "input": "ambient electronic pad"}
|
||||
body.update(overrides)
|
||||
openai_client.send_audio_generate_http_request(
|
||||
online_client.send_audio_generate_http_request(
|
||||
{"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
|
||||
def test_audio_generate_missing_input(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
|
||||
openai_client.send_audio_generate_http_request(
|
||||
def test_audio_generate_missing_input(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
|
||||
online_client.send_audio_generate_http_request(
|
||||
{
|
||||
"json": {"model": omni_server.model},
|
||||
"timeout": 120,
|
||||
|
||||
@@ -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
|
||||
"""HTTP + WebSocket validation for Qwen3-TTS Base, Higgs Audio V3 and Audex TTA ``/v1/audio/speech`` (and related routes)."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -13,8 +13,8 @@ from typing import Any
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.media import load_test_audio_data_url
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
|
||||
from tests.helpers.media import get_asset_path
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
|
||||
# Speech / voice-upload numeric caps for these tests only (not exported from vllm_omni).
|
||||
@@ -26,7 +26,7 @@ _SPEECH_API_MAX_NEW_TOKENS = 4096
|
||||
_VOICE_UPLOAD_MAX_CONSENT_LEN = 1024
|
||||
_VOICE_UPLOAD_MAX_REF_TEXT_CHARS = 8192
|
||||
_VOICE_UPLOAD_MAX_SPEAKER_DESCRIPTION_CHARS = 2048
|
||||
REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
|
||||
REF_AUDIO_URL = get_asset_path("qwen3_tts/clone_2.wav", as_data_url=True)
|
||||
REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
|
||||
# Qwen3-TTS 0.6B expects 1024-dim speaker embeddings (see serving_speech / talker hidden_size).
|
||||
_QWEN3_TTS_06B_SPEAKER_EMBEDDING_DIM = 1024
|
||||
@@ -57,6 +57,7 @@ _QWEN3_TTS_SPEECH = [
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
_HIGGS_AUDIO_V3_SPEECH = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
@@ -84,7 +85,7 @@ _AUDEX_TTA_SPEECH = [
|
||||
]
|
||||
|
||||
|
||||
def _apply_batch_overrides(body: dict[str, object], *, loc: str, overrides: dict[str, object]) -> None:
|
||||
def _apply_batch_overrides(body: dict[str, Any], *, loc: str, overrides: dict[str, object]) -> None:
|
||||
if loc == "batch":
|
||||
body.update(overrides)
|
||||
return
|
||||
@@ -112,8 +113,8 @@ def _pcm_wav_mono_bytes(*, duration_s: float = 1.1, sample_rate: int = 44100) ->
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
|
||||
def test_speech_malformed_json(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
|
||||
openai_client.send_audio_speech_http_request(
|
||||
def test_speech_malformed_json(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
|
||||
online_client.send_audio_speech_http_request(
|
||||
{
|
||||
"raw_body": "{",
|
||||
"timeout": 120,
|
||||
@@ -124,8 +125,8 @@ def test_speech_malformed_json(omni_server: OmniServer, openai_client: OpenAICli
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
|
||||
def test_speech_missing_required_fields(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
|
||||
openai_client.send_audio_speech_http_request(
|
||||
def test_speech_missing_required_fields(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
|
||||
online_client.send_audio_speech_http_request(
|
||||
{
|
||||
"json": {"model": omni_server.model},
|
||||
"timeout": 120,
|
||||
@@ -156,8 +157,8 @@ def test_speech_missing_required_fields(omni_server: OmniServer, openai_client:
|
||||
"ref_audio": None,
|
||||
"ref_text": None,
|
||||
},
|
||||
("CustomVoice", "does not support"),
|
||||
id="customvoice_task_not_supported",
|
||||
("Invalid voice", "vivian"),
|
||||
id="customvoice_invalid_voice",
|
||||
),
|
||||
pytest.param(
|
||||
{"task_type": "InvalidEnum"}, ("task_type", "literal_error", "CustomVoice"), id="task_type_invalid"
|
||||
@@ -180,13 +181,17 @@ def test_speech_missing_required_fields(omni_server: OmniServer, openai_client:
|
||||
("mutually exclusive", "speaker_embedding", "ref_audio"),
|
||||
id="speaker_embedding_with_ref_audio_mutually_exclusive",
|
||||
),
|
||||
pytest.param({"max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min"),
|
||||
pytest.param({"max_new_tokens": 0}, ("max_new_tokens", "greater_than_equal"), id="max_new_tokens_below_min"),
|
||||
pytest.param(
|
||||
{"max_new_tokens": _SPEECH_API_MAX_NEW_TOKENS + 1},
|
||||
("max_new_tokens", "exceed 4096"),
|
||||
id="max_new_tokens_above_max",
|
||||
),
|
||||
pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative"),
|
||||
pytest.param(
|
||||
{"seed": -(2**63) - 1},
|
||||
("seed", "greater_than_equal", "-9223372036854775808"),
|
||||
id="seed_below_min",
|
||||
),
|
||||
pytest.param({"seed": 2**63}, ("seed", "less_than_equal", "9223372036854775807"), id="seed_above_max"),
|
||||
pytest.param(
|
||||
{"initial_codec_chunk_frames": -1},
|
||||
@@ -242,13 +247,13 @@ def test_speech_missing_required_fields(omni_server: OmniServer, openai_client:
|
||||
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
|
||||
def test_speech_invalid_field_values(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
online_client: OnlineOmniClient,
|
||||
overrides: dict[str, object],
|
||||
err_message: str | tuple[str, ...],
|
||||
) -> None:
|
||||
ov = dict(overrides)
|
||||
embedding_only = bool(ov.pop(_SPEECH_INVALID_EMBEDDING_ONLY_BODY, False))
|
||||
body = {
|
||||
body: dict[str, Any] = {
|
||||
"model": omni_server.model,
|
||||
"input": "Hello.",
|
||||
"voice": "clone",
|
||||
@@ -261,7 +266,7 @@ def test_speech_invalid_field_values(
|
||||
body.pop("ref_audio", None)
|
||||
body.pop("ref_text", None)
|
||||
body.update(ov)
|
||||
openai_client.send_audio_speech_http_request(
|
||||
online_client.send_audio_speech_http_request(
|
||||
{"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
|
||||
)
|
||||
|
||||
@@ -274,7 +279,9 @@ def test_speech_invalid_field_values(
|
||||
pytest.param({"input": "", "response_format": "wav"}, ("input", "empty"), id="input_empty"),
|
||||
pytest.param({"input": " ", "response_format": "wav"}, ("input", "empty"), id="input_whitespace_only"),
|
||||
pytest.param(
|
||||
{"input": "Hello world.", "max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min"
|
||||
{"input": "Hello world.", "max_new_tokens": 0},
|
||||
("max_new_tokens", "greater_than_equal"),
|
||||
id="max_new_tokens_below_min",
|
||||
),
|
||||
pytest.param(
|
||||
{
|
||||
@@ -328,12 +335,12 @@ def test_speech_invalid_field_values(
|
||||
@pytest.mark.parametrize("omni_server", _HIGGS_AUDIO_V3_SPEECH, indirect=True)
|
||||
def test_speech_higgs_invalid_field_values(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
online_client: OnlineOmniClient,
|
||||
extra_json: dict[str, object],
|
||||
err_message: str | tuple[str, ...],
|
||||
) -> None:
|
||||
body = {"model": omni_server.model, **extra_json}
|
||||
openai_client.send_audio_speech_http_request(
|
||||
online_client.send_audio_speech_http_request(
|
||||
{"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
|
||||
)
|
||||
|
||||
@@ -359,12 +366,12 @@ def test_speech_higgs_invalid_field_values(
|
||||
@pytest.mark.parametrize("omni_server", _AUDEX_TTA_SPEECH, indirect=True)
|
||||
def test_speech_audex_tta_invalid_field_values(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
online_client: OnlineOmniClient,
|
||||
overrides: dict[str, object],
|
||||
err_message: str | tuple[str, ...],
|
||||
) -> None:
|
||||
body = {"model": omni_server.model, "input": _AUDEX_TTA_CAPTION, **overrides}
|
||||
openai_client.send_audio_speech_http_request(
|
||||
online_client.send_audio_speech_http_request(
|
||||
{"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
|
||||
)
|
||||
|
||||
@@ -375,15 +382,15 @@ def test_speech_audex_tta_invalid_field_values(
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
|
||||
def test_speech_batch_empty_items(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
|
||||
openai_client.send_audio_speech_batch_http_request(
|
||||
def test_speech_batch_empty_items(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
|
||||
online_client.send_audio_speech_batch_http_request(
|
||||
{"json": {"items": []}, "timeout": 120, "err_code": (400, 422), "err_message": ("items", "least 1 item")}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
|
||||
def test_speech_batch_missing_items(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
|
||||
openai_client.send_audio_speech_batch_http_request(
|
||||
def test_speech_batch_missing_items(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
|
||||
online_client.send_audio_speech_batch_http_request(
|
||||
{
|
||||
"json": {"model": omni_server.model},
|
||||
"timeout": 120,
|
||||
@@ -553,14 +560,14 @@ def test_speech_batch_missing_items(omni_server: OmniServer, openai_client: Open
|
||||
pytest.param(
|
||||
"batch",
|
||||
{"max_new_tokens": 0},
|
||||
("max_new_tokens", "least 1"),
|
||||
("max_new_tokens", "greater_than_equal"),
|
||||
id="batch_max_new_tokens_below_min",
|
||||
marks=_SKIP_ISSUE_3649,
|
||||
),
|
||||
pytest.param(
|
||||
"item",
|
||||
{"max_new_tokens": 0},
|
||||
("max_new_tokens", "least 1"),
|
||||
("max_new_tokens", "greater_than_equal"),
|
||||
id="item_max_new_tokens_below_min",
|
||||
marks=_SKIP_ISSUE_3649,
|
||||
),
|
||||
@@ -618,14 +625,14 @@ def test_speech_batch_missing_items(omni_server: OmniServer, openai_client: Open
|
||||
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
|
||||
def test_speech_batch_invalid_field_values(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
online_client: OnlineOmniClient,
|
||||
loc: str,
|
||||
overrides: dict[str, object],
|
||||
err_message: str | tuple[str, ...],
|
||||
) -> None:
|
||||
ov = dict(overrides)
|
||||
strip_clone = bool(ov.pop(_SPEECH_BATCH_NO_CLONE_FIELDS, False))
|
||||
body = {
|
||||
body: dict[str, Any] = {
|
||||
"model": omni_server.model,
|
||||
"voice": "clone",
|
||||
"ref_audio": REF_AUDIO_URL,
|
||||
@@ -638,7 +645,7 @@ def test_speech_batch_invalid_field_values(
|
||||
body.pop("ref_audio", None)
|
||||
body.pop("ref_text", None)
|
||||
_apply_batch_overrides(body, loc=loc, overrides=ov)
|
||||
openai_client.send_audio_speech_batch_http_request(
|
||||
online_client.send_audio_speech_batch_http_request(
|
||||
{"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
|
||||
)
|
||||
|
||||
@@ -670,7 +677,7 @@ _SPEECH_STREAM_WS_SESSION_STREAM_AUDIO_WAV_FRAMES = object()
|
||||
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
|
||||
def test_speech_stream_invalid_requests(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
online_client: OnlineOmniClient,
|
||||
send_frames_spec: Any,
|
||||
err_message: str | tuple[str, ...],
|
||||
) -> None:
|
||||
@@ -687,7 +694,7 @@ def test_speech_stream_invalid_requests(
|
||||
else:
|
||||
send_frames = send_frames_spec
|
||||
assert isinstance(send_frames, str)
|
||||
openai_client.send_audio_speech_stream_ws_request(
|
||||
online_client.send_audio_speech_stream_ws_request(
|
||||
{
|
||||
"send_frames": send_frames,
|
||||
"timeout": 120,
|
||||
@@ -893,7 +900,7 @@ def _voices_upload_multipart_files(kind: str | None) -> dict[str, Any] | None:
|
||||
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
|
||||
def test_voices_create_invalid_requests(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
online_client: OnlineOmniClient,
|
||||
multipart_kind: str | None,
|
||||
form_template: dict[str, Any],
|
||||
err_message: str | tuple[str, ...],
|
||||
@@ -909,7 +916,7 @@ def test_voices_create_invalid_requests(
|
||||
files = _voices_upload_multipart_files(multipart_kind)
|
||||
if files is not None:
|
||||
cfg["files"] = files
|
||||
openai_client.send_audio_voices_create_http_request(cfg)
|
||||
online_client.send_audio_voices_create_http_request(cfg)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -962,13 +969,13 @@ def test_voices_create_invalid_requests(
|
||||
@pytest.mark.parametrize("omni_server", _QWEN3_TTS_SPEECH, indirect=True)
|
||||
def test_voices_delete_invalid_requests(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
online_client: OnlineOmniClient,
|
||||
voice_name: str,
|
||||
err_code: int | tuple[int, ...],
|
||||
err_message: str | tuple[str, ...],
|
||||
) -> None:
|
||||
"""DELETE ``/v1/audio/voices/{name}``: missing voice (404), odd segments (often 404), ``.``paths (405)."""
|
||||
openai_client.send_audio_voices_delete_http_request(
|
||||
online_client.send_audio_voices_delete_http_request(
|
||||
{
|
||||
"name": voice_name,
|
||||
"timeout": 120,
|
||||
@@ -976,3 +983,68 @@ def test_voices_delete_invalid_requests(
|
||||
"err_message": err_message,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
# ─── POST /v1/audio/speech · MiniMax Music 3 ───
|
||||
|
||||
# Text-to-music on the speech endpoint. Most of the speech contract does not
|
||||
# apply: there is no speaker, no reference audio and no temperature, and the
|
||||
# lyrics alone are not a request. Each case below pins a rejection so a caller
|
||||
# who assumes the usual TTS shape finds out immediately.
|
||||
_MINIMAX_MUSIC3_SPEECH = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model="MiniMaxAI/MiniMax-Music3",
|
||||
stage_config_path=get_deploy_config_path("minimax_music3.yaml"),
|
||||
server_args=_SPEECH_SERVER_ARGS,
|
||||
),
|
||||
id="minimax_music3",
|
||||
marks=hardware_marks(res={"cuda": "H100"}),
|
||||
),
|
||||
]
|
||||
|
||||
_MINIMAX_MUSIC3_LYRICS = "[Verse]\nWalking down the empty street at midnight"
|
||||
_MINIMAX_MUSIC3_CAPTION = "A melancholic lo-fi hip-hop track at 85 BPM in F minor."
|
||||
# 25 frames per second, capped at six minutes of song.
|
||||
_MINIMAX_MUSIC3_MAX_FRAMES = 9000
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"overrides, err_message",
|
||||
[
|
||||
pytest.param({"instructions": None}, ("instructions",), id="caption_missing"),
|
||||
pytest.param({"instructions": " "}, ("instructions",), id="caption_blank"),
|
||||
pytest.param({"input": " "}, ("input",), id="lyrics_blank"),
|
||||
pytest.param({"voice": "alloy"}, ("voice",), id="voice_rejected"),
|
||||
pytest.param({"speed": 1.5}, ("speed",), id="speed_rejected"),
|
||||
pytest.param(
|
||||
{"max_new_tokens": _MINIMAX_MUSIC3_MAX_FRAMES + 1},
|
||||
("max_new_tokens", str(_MINIMAX_MUSIC3_MAX_FRAMES)),
|
||||
id="max_new_tokens_over_cap",
|
||||
),
|
||||
pytest.param({"max_new_tokens": 0}, ("max_new_tokens",), id="max_new_tokens_zero"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("omni_server", _MINIMAX_MUSIC3_SPEECH, indirect=True)
|
||||
def test_minimax_music3_speech_invalid_field_values(
|
||||
omni_server: OmniServer,
|
||||
online_client: OnlineOmniClient,
|
||||
overrides: dict[str, object],
|
||||
err_message: str | tuple[str, ...],
|
||||
) -> None:
|
||||
body: dict[str, Any] = {
|
||||
"model": omni_server.model,
|
||||
"input": _MINIMAX_MUSIC3_LYRICS,
|
||||
"instructions": _MINIMAX_MUSIC3_CAPTION,
|
||||
"seed": 1,
|
||||
"max_new_tokens": 250,
|
||||
"response_format": "wav",
|
||||
}
|
||||
for key, value in overrides.items():
|
||||
if value is None:
|
||||
body.pop(key, None)
|
||||
else:
|
||||
body[key] = value
|
||||
online_client.send_audio_speech_http_request(
|
||||
{"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
|
||||
)
|
||||
@@ -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
|
||||
"""``POST /v1/images/edits`` multipart validation (live ``Qwen/Qwen-Image-Edit``)."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -9,7 +9,7 @@ from typing import Any
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
|
||||
|
||||
pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
|
||||
|
||||
@@ -72,7 +72,7 @@ def _finalize_edits_form_data(template: dict[str, Any], server_model: str) -> di
|
||||
@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
|
||||
def test_images_edits_invalid_requests(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
online_client: OnlineOmniClient,
|
||||
tiny_png_bytes: bytes,
|
||||
include_image: bool,
|
||||
data_template: dict[str, Any],
|
||||
@@ -87,4 +87,4 @@ def test_images_edits_invalid_requests(
|
||||
}
|
||||
if include_image:
|
||||
cfg["files"] = {"image": ("x.png", tiny_png_bytes, "image/png")}
|
||||
openai_client.send_images_edits_http_request(cfg)
|
||||
online_client.send_images_edits_http_request(cfg)
|
||||
@@ -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
|
||||
"""``POST /v1/images/generations`` validation (live ``Qwen/Qwen-Image``)."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -9,7 +9,7 @@ from typing import Any
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
|
||||
|
||||
pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
|
||||
|
||||
@@ -126,7 +126,7 @@ def _minimal_images_gen_json(omni_server: OmniServer) -> dict[str, object]:
|
||||
@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
|
||||
def test_images_generations_invalid_requests(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
online_client: OnlineOmniClient,
|
||||
body_spec: Any,
|
||||
err_message: str | tuple[str, ...],
|
||||
) -> None:
|
||||
@@ -137,6 +137,6 @@ def test_images_generations_invalid_requests(
|
||||
assert isinstance(body_spec, dict)
|
||||
body = _minimal_images_gen_json(omni_server)
|
||||
body.update(body_spec)
|
||||
openai_client.send_images_generations_http_request(
|
||||
online_client.send_images_generations_http_request(
|
||||
{"json": body, "timeout": 300, "err_code": 400, "err_message": err_message}
|
||||
)
|
||||
@@ -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
|
||||
"""Invalid inputs on Qwen3-Omni: ``POST /v1/chat/completions``, ``WS /v1/video/chat/stream``, ``WS /v1/realtime``."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -10,7 +10,7 @@ from typing import Any
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
|
||||
pytestmark = [pytest.mark.slow, pytest.mark.omni]
|
||||
@@ -72,6 +72,9 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
|
||||
body["top_logprobs"] = 5
|
||||
elif case_id == "speaker_unknown":
|
||||
body["speaker"] = "zz_invalid_qwen3_omni_chat_speaker_xyz"
|
||||
elif case_id == "audio_format_unsupported":
|
||||
body["modalities"] = ["text", "audio"]
|
||||
body["audio"] = {"voice": "alloy", "format": "aac"}
|
||||
else:
|
||||
raise AssertionError(f"unknown chat completions invalid case_id {case_id!r}")
|
||||
return {"json": body, "timeout": 120}
|
||||
@@ -119,7 +122,7 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
|
||||
pytest.param(
|
||||
"response_format_json_schema_incomplete",
|
||||
400,
|
||||
("response_format", "value_error", "json_schema"),
|
||||
("response_format", "BadRequestError", "json_schema"),
|
||||
id="invalid_response_format_json_schema",
|
||||
),
|
||||
pytest.param("logprobs_wrong_type", 400, "logprobs", id="logprobs_wrong_type", marks=_SKIP_ISSUE_3649),
|
||||
@@ -135,12 +138,18 @@ def _chat_completions_request_without_expectations(omni_server: OmniServer, case
|
||||
("Invalid speaker", "Supported"),
|
||||
id="speaker_unknown_preset",
|
||||
),
|
||||
pytest.param(
|
||||
"audio_format_unsupported",
|
||||
400,
|
||||
("Invalid audio format", "aac", "Supported formats"),
|
||||
id="unsupported_audio_format",
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
|
||||
def test_chat_completions_invalid_requests(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
online_client: OnlineOmniClient,
|
||||
case_id: str,
|
||||
err_code: int | tuple[int, ...],
|
||||
err_message: str | tuple[str, ...],
|
||||
@@ -148,7 +157,7 @@ def test_chat_completions_invalid_requests(
|
||||
cfg = _chat_completions_request_without_expectations(omni_server, case_id)
|
||||
cfg["err_code"] = err_code
|
||||
cfg["err_message"] = err_message
|
||||
openai_client.send_chat_completions_http_request(cfg)[0]
|
||||
online_client.send_chat_completions_http_request(cfg)[0]
|
||||
|
||||
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
@@ -174,7 +183,7 @@ _VIDEO_CHAT_WS_SESSION_MODALITIES_SCALAR = object()
|
||||
@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
|
||||
def test_video_chat_stream_invalid_requests(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
online_client: OnlineOmniClient,
|
||||
send_frames_spec: Any,
|
||||
err_message: str,
|
||||
) -> None:
|
||||
@@ -183,7 +192,7 @@ def test_video_chat_stream_invalid_requests(
|
||||
else:
|
||||
send_frames = send_frames_spec
|
||||
assert isinstance(send_frames, str)
|
||||
openai_client.send_video_chat_stream_ws_request(
|
||||
online_client.send_video_chat_stream_ws_request(
|
||||
{
|
||||
"send_frames": send_frames,
|
||||
"timeout": 120,
|
||||
@@ -243,11 +252,12 @@ _REALTIME_WS_INVALID_AUDIO_APPEND = object()
|
||||
@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
|
||||
def test_realtime_invalid_requests(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
online_client: OnlineOmniClient,
|
||||
send_frames_spec: Any,
|
||||
ws_error_code: str,
|
||||
err_message: str | tuple[str, ...],
|
||||
) -> None:
|
||||
send_frames: Any
|
||||
if send_frames_spec is _REALTIME_WS_MODEL_MISMATCH:
|
||||
send_frames = json.dumps(
|
||||
{
|
||||
@@ -262,7 +272,7 @@ def test_realtime_invalid_requests(
|
||||
]
|
||||
else:
|
||||
send_frames = send_frames_spec
|
||||
openai_client.send_realtime_ws_request(
|
||||
online_client.send_realtime_ws_request(
|
||||
{
|
||||
"send_frames": send_frames,
|
||||
"timeout": 120,
|
||||
|
||||
@@ -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
|
||||
"""Omni sleep/wakeup JSON validation and level-2 wake contract (live ``Qwen/Qwen-Image``)."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -9,7 +9,7 @@ from typing import Any, Literal
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
|
||||
|
||||
pytestmark = [pytest.mark.slow, pytest.mark.diffusion, pytest.mark.omni]
|
||||
|
||||
@@ -86,7 +86,7 @@ _QWEN_IMAGE = [
|
||||
@pytest.mark.parametrize("omni_server", _QWEN_IMAGE, indirect=True)
|
||||
def test_omni_sleep_wakeup_invalid_requests(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
online_client: OnlineOmniClient,
|
||||
endpoint: Literal["sleep", "wakeup"],
|
||||
json_body: dict[str, Any],
|
||||
err_message: str | tuple[str, ...],
|
||||
@@ -98,9 +98,9 @@ def test_omni_sleep_wakeup_invalid_requests(
|
||||
"err_message": err_message,
|
||||
}
|
||||
if endpoint == "sleep":
|
||||
openai_client.send_omni_sleep_http_request(cfg)
|
||||
online_client.send_omni_sleep_http_request(cfg)
|
||||
else:
|
||||
openai_client.send_omni_wakeup_http_request(cfg)
|
||||
online_client.send_omni_wakeup_http_request(cfg)
|
||||
|
||||
|
||||
# Keep this after the invalid-JSON cases: sleep(level=2) discards weights and the
|
||||
@@ -108,7 +108,7 @@ def test_omni_sleep_wakeup_invalid_requests(
|
||||
@pytest.mark.parametrize("omni_server", _QWEN_IMAGE, indirect=True)
|
||||
def test_wakeup_after_level2_sleep_fails(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
online_client: OnlineOmniClient,
|
||||
) -> None:
|
||||
"""Regression for #4473 Repro A via HTTP: wake after sleep(level=2) → 501.
|
||||
|
||||
@@ -116,7 +116,7 @@ def test_wakeup_after_level2_sleep_fails(
|
||||
Live serve surfaces ``NotImplementedError`` as OpenAI-style JSON with code 501.
|
||||
"""
|
||||
assert omni_server is not None
|
||||
sleep_resps = openai_client.send_omni_sleep_http_request(
|
||||
sleep_resps = online_client.send_omni_sleep_http_request(
|
||||
{
|
||||
"json": {"stage_ids": [0], "level": 2},
|
||||
"timeout": 180,
|
||||
@@ -124,10 +124,11 @@ def test_wakeup_after_level2_sleep_fails(
|
||||
)
|
||||
sleep_resp = sleep_resps[0]
|
||||
assert sleep_resp.status_code == 200, sleep_resp
|
||||
assert sleep_resp.json_body is not None
|
||||
assert sleep_resp.json_body.get("status") == "SUCCESS", sleep_resp.json_body
|
||||
json_body = sleep_resp.json_body
|
||||
assert isinstance(json_body, dict)
|
||||
assert json_body.get("status") == "SUCCESS", json_body
|
||||
|
||||
openai_client.send_omni_wakeup_http_request(
|
||||
online_client.send_omni_wakeup_http_request(
|
||||
{
|
||||
"json": {"stage_ids": [0]},
|
||||
"timeout": 120,
|
||||
|
||||
@@ -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
|
||||
"""HTTP video job routes (live ``Wan-AI/Wan2.2-T2V-A14B-Diffusers``).
|
||||
|
||||
WebSocket ``/v1/video/chat/stream`` and ``/v1/realtime`` invalid-input checks live in
|
||||
@@ -13,7 +13,7 @@ from typing import Any
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler
|
||||
from tests.helpers.runtime import OmniServer, OmniServerParams, OnlineOmniClient
|
||||
|
||||
pytestmark = [pytest.mark.slow, pytest.mark.diffusion]
|
||||
|
||||
@@ -54,12 +54,12 @@ def _minimal_video_form_data() -> dict[str, str]:
|
||||
@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
|
||||
def test_post_videos_invalid_requests(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
online_client: OnlineOmniClient,
|
||||
form_extra: dict[str, str] | None,
|
||||
err_message: str | tuple[str, ...],
|
||||
) -> None:
|
||||
data: dict[str, str] = {} if form_extra is None else {**_minimal_video_form_data(), **form_extra}
|
||||
openai_client.send_videos_create_http_request(
|
||||
online_client.send_videos_create_http_request(
|
||||
{"data": data, "timeout": 120, "err_code": 400, "err_message": err_message}
|
||||
)
|
||||
|
||||
@@ -85,12 +85,12 @@ def test_post_videos_invalid_requests(
|
||||
@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
|
||||
def test_post_videos_sync_invalid_requests(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
online_client: OnlineOmniClient,
|
||||
form_extra: dict[str, str] | None,
|
||||
err_message: str | tuple[str, ...],
|
||||
) -> None:
|
||||
data: dict[str, str] = {} if form_extra is None else {**_minimal_video_form_data(), **form_extra}
|
||||
openai_client.send_videos_sync_http_request(
|
||||
online_client.send_videos_sync_http_request(
|
||||
{"data": data, "timeout": 120, "err_code": 400, "err_message": err_message}
|
||||
)
|
||||
|
||||
@@ -111,11 +111,11 @@ def test_post_videos_sync_invalid_requests(
|
||||
@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
|
||||
def test_get_videos_list_invalid_requests(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
online_client: OnlineOmniClient,
|
||||
params: dict[str, Any],
|
||||
err_message: str | tuple[str, ...],
|
||||
) -> None:
|
||||
openai_client.send_videos_list_http_request(
|
||||
online_client.send_videos_list_http_request(
|
||||
{"params": params, "timeout": 120, "err_code": 400, "err_message": err_message}
|
||||
)
|
||||
|
||||
@@ -126,8 +126,8 @@ def test_get_videos_list_invalid_requests(
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
|
||||
def test_get_video_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
|
||||
openai_client.send_video_retrieve_http_request(
|
||||
def test_get_video_invalid_requests(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
|
||||
online_client.send_video_retrieve_http_request(
|
||||
{
|
||||
"video_id": "does-not-exist",
|
||||
"timeout": 120,
|
||||
@@ -143,8 +143,8 @@ def test_get_video_invalid_requests(omni_server: OmniServer, openai_client: Open
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
|
||||
def test_delete_video_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
|
||||
openai_client.send_video_delete_http_request(
|
||||
def test_delete_video_invalid_requests(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
|
||||
online_client.send_video_delete_http_request(
|
||||
{
|
||||
"video_id": "does-not-exist",
|
||||
"timeout": 120,
|
||||
@@ -160,8 +160,8 @@ def test_delete_video_invalid_requests(omni_server: OmniServer, openai_client: O
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", _WAN_T2V, indirect=True)
|
||||
def test_get_video_content_invalid_requests(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None:
|
||||
openai_client.send_video_content_http_request(
|
||||
def test_get_video_content_invalid_requests(omni_server: OmniServer, online_client: OnlineOmniClient) -> None:
|
||||
online_client.send_video_content_http_request(
|
||||
{
|
||||
"video_id": "does-not-exist",
|
||||
"timeout": 120,
|
||||
|
||||
@@ -1,3 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""
|
||||
HunyuanImage-3.0-Instruct reliability integration tests (DiT-only deploy).
|
||||
|
||||
@@ -35,7 +38,7 @@ from tests.dfx.reliability.helpers import (
|
||||
worker_residual_timeout_after_kill_signal,
|
||||
)
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OpenAIClientHandler
|
||||
from tests.helpers.runtime import OnlineOmniClient
|
||||
|
||||
RELIABILITY_SCENARIOS: list[dict[str, Any]] = [
|
||||
{
|
||||
@@ -133,8 +136,8 @@ def _image_request_config(omni_server: Any) -> dict[str, Any]:
|
||||
}
|
||||
|
||||
|
||||
def _submit_image_generation(openai_client: OpenAIClientHandler, omni_server: Any) -> None:
|
||||
responses = openai_client.send_images_generations_http_request(_image_request_config(omni_server))
|
||||
def _submit_image_generation(online_client: OnlineOmniClient, omni_server: Any) -> None:
|
||||
responses = online_client.send_images_generations_http_request(_image_request_config(omni_server))
|
||||
resp = responses[0]
|
||||
if resp.status_code >= 400:
|
||||
raise RuntimeError(
|
||||
@@ -189,7 +192,7 @@ def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -
|
||||
@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True)
|
||||
def test_reliability_fault_gpu_oom_image_large_request_failure(
|
||||
omni_server_function,
|
||||
openai_client_function,
|
||||
online_client_function,
|
||||
) -> None:
|
||||
stage_config_path = getattr(omni_server_function, "stage_config_path", None)
|
||||
device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
|
||||
@@ -202,7 +205,7 @@ def test_reliability_fault_gpu_oom_image_large_request_failure(
|
||||
)
|
||||
try:
|
||||
try:
|
||||
_submit_image_generation(openai_client_function, omni_server_function)
|
||||
_submit_image_generation(online_client_function, omni_server_function)
|
||||
except Exception as exc:
|
||||
assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
|
||||
else:
|
||||
@@ -222,10 +225,10 @@ def test_reliability_fault_gpu_oom_image_large_request_failure(
|
||||
@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_request_failure(
|
||||
omni_server_after_fault_function,
|
||||
openai_client_function,
|
||||
online_client_function,
|
||||
) -> None:
|
||||
try:
|
||||
_submit_image_generation(openai_client_function, omni_server_after_fault_function)
|
||||
_submit_image_generation(online_client_function, omni_server_after_fault_function)
|
||||
except Exception as exc:
|
||||
assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
|
||||
else:
|
||||
@@ -244,7 +247,7 @@ def test_reliability_fault_process_kill_request_failure(
|
||||
@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
|
||||
omni_server_after_fault_function,
|
||||
openai_client_function,
|
||||
online_client_function,
|
||||
) -> None:
|
||||
host = omni_server_after_fault_function.host
|
||||
port = omni_server_after_fault_function.port
|
||||
@@ -280,7 +283,7 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
|
||||
start = time.monotonic()
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
|
||||
futures = [
|
||||
executor.submit(_submit_image_generation, openai_client_function, omni_server_after_fault_function)
|
||||
executor.submit(_submit_image_generation, online_client_function, omni_server_after_fault_function)
|
||||
for _ in range(3)
|
||||
]
|
||||
done, pending = concurrent.futures.wait(
|
||||
@@ -335,12 +338,12 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
|
||||
@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_worker_with_load_request_failure(
|
||||
omni_server_function,
|
||||
openai_client_function,
|
||||
online_client_function,
|
||||
fault_injector: FaultInjector,
|
||||
) -> None:
|
||||
scenario = "kill_worker_with_load"
|
||||
load_result = run_fault_injection_with_rate_load(
|
||||
submit_request=lambda: _submit_image_generation(openai_client_function, omni_server_function),
|
||||
submit_request=lambda: _submit_image_generation(online_client_function, omni_server_function),
|
||||
inject_fault=lambda: fault_injector(omni_server_function),
|
||||
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
|
||||
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
|
||||
@@ -363,13 +366,13 @@ def test_reliability_fault_process_kill_worker_with_load_request_failure(
|
||||
@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup(
|
||||
omni_server_function,
|
||||
openai_client_function,
|
||||
online_client_function,
|
||||
signal_name: str,
|
||||
) -> None:
|
||||
scenario = f"kill_serve_root_with_load_{signal_name.lower()}"
|
||||
injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
|
||||
load_result = run_fault_injection_with_rate_load(
|
||||
submit_request=lambda: _submit_image_generation(openai_client_function, omni_server_function),
|
||||
submit_request=lambda: _submit_image_generation(online_client_function, omni_server_function),
|
||||
inject_fault=lambda: injector(omni_server_function),
|
||||
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
|
||||
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
|
||||
@@ -448,7 +451,7 @@ def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup(
|
||||
@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
|
||||
omni_server_function,
|
||||
openai_client_function,
|
||||
online_client_function,
|
||||
signal_name: str,
|
||||
) -> None:
|
||||
scenario = f"kill_serve_tree_with_load_{signal_name.lower()}"
|
||||
@@ -458,7 +461,7 @@ def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
|
||||
inter_kill_wait_seconds=0.1,
|
||||
)
|
||||
load_result = run_fault_injection_with_rate_load(
|
||||
submit_request=lambda: _submit_image_generation(openai_client_function, omni_server_function),
|
||||
submit_request=lambda: _submit_image_generation(online_client_function, omni_server_function),
|
||||
inject_fault=lambda: injector(omni_server_function),
|
||||
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
|
||||
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
|
||||
|
||||
@@ -1,3 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""
|
||||
Qwen3-Omni reliability integration tests.
|
||||
"""
|
||||
@@ -63,7 +66,7 @@ DEPLOY_CONFIGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "vll
|
||||
|
||||
def _default_oom_device_spec() -> str:
|
||||
"""Use currently visible CUDA ordinals to avoid invalid device index in sidecar."""
|
||||
count = torch.accelerator.device_count()
|
||||
count = torch.accelerator.device_count() or 0
|
||||
if count <= 0:
|
||||
return "0"
|
||||
return ",".join(str(i) for i in range(count))
|
||||
@@ -71,7 +74,7 @@ def _default_oom_device_spec() -> str:
|
||||
|
||||
OOM_INJECTION_CONFIG = {
|
||||
"device": _default_oom_device_spec(),
|
||||
"target_mem_ratio": 0.92,
|
||||
"target_mem_ratio": 0.95,
|
||||
"hold_seconds": 0,
|
||||
"startup_timeout_sec": 20,
|
||||
"strict": False,
|
||||
@@ -414,7 +417,7 @@ def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech(
|
||||
@pytest.mark.slow
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
|
||||
def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_function, openai_client_function) -> None:
|
||||
def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_function, online_client_function) -> None:
|
||||
device_spec = resolve_oom_device_spec(
|
||||
OOM_INJECTION_CONFIG,
|
||||
_deploy_config_path_from_omni_server(omni_server_function),
|
||||
@@ -444,7 +447,7 @@ def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_functi
|
||||
"key_words": {"audio": ["test"]},
|
||||
}
|
||||
try:
|
||||
openai_client_function.send_omni_request(request_config, request_num=1)
|
||||
online_client_function.send_omni_request(request_config, request_num=1)
|
||||
except Exception as exc:
|
||||
assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
|
||||
else:
|
||||
@@ -456,7 +459,7 @@ def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_functi
|
||||
@pytest.mark.slow
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
|
||||
def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_function, openai_client_function) -> None:
|
||||
def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_function, online_client_function) -> None:
|
||||
device_spec = resolve_oom_device_spec(
|
||||
OOM_INJECTION_CONFIG,
|
||||
_deploy_config_path_from_omni_server(omni_server_function),
|
||||
@@ -487,7 +490,7 @@ def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_funct
|
||||
"key_words": {"audio": ["test"]},
|
||||
}
|
||||
try:
|
||||
openai_client_function.send_omni_request(request_config, request_num=4)
|
||||
online_client_function.send_omni_request(request_config, request_num=4)
|
||||
except Exception as exc:
|
||||
assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
|
||||
else:
|
||||
@@ -505,7 +508,7 @@ def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_funct
|
||||
)
|
||||
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_request_failure(
|
||||
omni_server_after_fault_function, openai_client_function
|
||||
omni_server_after_fault_function, online_client_function
|
||||
) -> None:
|
||||
messages = dummy_messages_from_mix_data(
|
||||
system_prompt=_get_system_prompt(),
|
||||
@@ -519,7 +522,7 @@ def test_reliability_fault_process_kill_request_failure(
|
||||
"key_words": {"text": ["beijing"]},
|
||||
}
|
||||
try:
|
||||
openai_client_function.send_omni_request(request_config, request_num=1)
|
||||
online_client_function.send_omni_request(request_config, request_num=1)
|
||||
except Exception as exc:
|
||||
assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
|
||||
else:
|
||||
@@ -567,7 +570,7 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
|
||||
assert elapsed <= _PROCESS_KILL_WORKER_CHAT_FF_MAX_ELAPSED_SEC + 1.0, (
|
||||
f"[process_kill fast_fail] request did not fail fast after fault: {elapsed:.2f}s"
|
||||
)
|
||||
assert ff_status >= 500, (
|
||||
assert ff_status is not None and ff_status >= 500, (
|
||||
f"[process_kill fast_fail] expected server-side failure after fault, "
|
||||
f"got status={ff_status}, body={ff_body[:200]!r}"
|
||||
)
|
||||
@@ -652,13 +655,13 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
|
||||
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_worker_with_load_request_failure(
|
||||
omni_server_function,
|
||||
openai_client_function,
|
||||
online_client_function,
|
||||
fault_injector: FaultInjector,
|
||||
) -> None:
|
||||
request_config = _chat_request_config(omni_server_function)
|
||||
scenario = "kill_worker_with_load"
|
||||
load_result = run_fault_injection_with_rate_load(
|
||||
submit_request=lambda: openai_client_function.send_omni_request(request_config, request_num=1),
|
||||
submit_request=lambda: online_client_function.send_omni_request(request_config, request_num=1),
|
||||
inject_fault=lambda: fault_injector(omni_server_function),
|
||||
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
|
||||
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
|
||||
@@ -704,14 +707,14 @@ def test_reliability_fault_process_kill_serve_root_no_load_fast_fail_and_cleanup
|
||||
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup(
|
||||
omni_server_function,
|
||||
openai_client_function,
|
||||
online_client_function,
|
||||
signal_name: str,
|
||||
) -> None:
|
||||
request_config = _chat_request_config(omni_server_function)
|
||||
scenario = f"kill_serve_root_with_load_{signal_name.lower()}"
|
||||
injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
|
||||
load_result = run_fault_injection_with_rate_load(
|
||||
submit_request=lambda: openai_client_function.send_omni_request(request_config, request_num=1),
|
||||
submit_request=lambda: online_client_function.send_omni_request(request_config, request_num=1),
|
||||
inject_fault=lambda: injector(omni_server_function),
|
||||
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
|
||||
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
|
||||
@@ -764,7 +767,7 @@ def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup(
|
||||
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
|
||||
omni_server_function,
|
||||
openai_client_function,
|
||||
online_client_function,
|
||||
signal_name: str,
|
||||
) -> None:
|
||||
request_config = _chat_request_config(omni_server_function)
|
||||
@@ -775,7 +778,7 @@ def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
|
||||
inter_kill_wait_seconds=0.1,
|
||||
)
|
||||
load_result = run_fault_injection_with_rate_load(
|
||||
submit_request=lambda: openai_client_function.send_omni_request(request_config, request_num=1),
|
||||
submit_request=lambda: online_client_function.send_omni_request(request_config, request_num=1),
|
||||
inject_fault=lambda: injector(omni_server_function),
|
||||
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
|
||||
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
|
||||
|
||||
@@ -1,508 +0,0 @@
|
||||
"""
|
||||
VoxCPM2 reliability integration tests.
|
||||
|
||||
Scenarios follow ``docs/user_guide/fault_injection_reliability_matrix.md``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import concurrent.futures
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.dfx.conftest import (
|
||||
assert_fault_exception,
|
||||
create_reliability_omni_server_params,
|
||||
resolve_oom_device_spec,
|
||||
)
|
||||
from tests.dfx.reliability.helpers import (
|
||||
FaultInjector,
|
||||
assert_no_server_tree_process_residual_and_gpu_release,
|
||||
extract_openai_error_contract_from_bytes,
|
||||
get_health_raw,
|
||||
inject_gpu_oom,
|
||||
make_process_kill_fault_injector,
|
||||
make_server_root_kill_fault_injector,
|
||||
make_server_tree_kill_fault_injector,
|
||||
post_json_raw,
|
||||
run_fault_injection_with_rate_load,
|
||||
stop_gpu_oom_hogs,
|
||||
worker_residual_timeout_after_kill_signal,
|
||||
)
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.media import load_test_audio_data_url
|
||||
from tests.helpers.runtime import OpenAIClientHandler
|
||||
|
||||
RELIABILITY_SCENARIOS: list[dict[str, Any]] = [
|
||||
{
|
||||
"test_name": "voxcpm2_reliability_default",
|
||||
"server_params": {
|
||||
"model": "openbmb/VoxCPM2",
|
||||
"stage_config_name": "voxcpm2.yaml",
|
||||
# VoxCPM2 uses a custom HF tokenizer; CLI default trust_remote_code=False
|
||||
# would otherwise override deploy YAML settings during server startup.
|
||||
"server_args": ["--trust-remote-code", "--disable-log-stats"],
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
DEPLOY_CONFIGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "vllm_omni" / "deploy"
|
||||
OOM_INJECTION_CONFIG = {
|
||||
"target_mem_ratio": 0.96,
|
||||
"hold_seconds": 0,
|
||||
"startup_timeout_sec": 20,
|
||||
"strict": False,
|
||||
}
|
||||
FAULT_ERROR_KEYWORDS = (
|
||||
"the request failed",
|
||||
"oom",
|
||||
"out of memory",
|
||||
"cuda",
|
||||
"orchestrator",
|
||||
"timeout",
|
||||
"connection",
|
||||
"500",
|
||||
"503",
|
||||
)
|
||||
PROCESS_KILL_ERROR_KEYWORDS = (
|
||||
"timeout",
|
||||
"did not complete within",
|
||||
"connection",
|
||||
"engine",
|
||||
"orchestrator",
|
||||
"dead",
|
||||
"internal",
|
||||
"500",
|
||||
"503",
|
||||
)
|
||||
RUNTIME_WORKER_PATTERN = "VLLM::"
|
||||
SERVE_SIGNAL_PARAMS = [
|
||||
pytest.param("SIGTERM", id="sigterm"),
|
||||
pytest.param("SIGINT", id="sigint"),
|
||||
pytest.param("SIGKILL", id="sigkill"),
|
||||
]
|
||||
SERVE_WITH_LOAD_SIGNAL_PARAMS = [
|
||||
pytest.param("SIGTERM", id="sigterm"),
|
||||
pytest.param("SIGINT", id="sigint"),
|
||||
pytest.param("SIGKILL", id="sigkill"),
|
||||
]
|
||||
TREE_SIGNAL_PARAMS = [
|
||||
pytest.param("SIGTERM", id="sigterm"),
|
||||
pytest.param("SIGKILL", id="sigkill"),
|
||||
]
|
||||
TREE_WITH_LOAD_SIGNAL_PARAMS = [
|
||||
pytest.param("SIGTERM", id="sigterm"),
|
||||
pytest.param("SIGKILL", id="sigkill"),
|
||||
]
|
||||
WORKER_SIGNAL_FAULT_PARAMS = [
|
||||
pytest.param(
|
||||
make_process_kill_fault_injector(
|
||||
grep_patterns=RUNTIME_WORKER_PATTERN,
|
||||
signal_name="SIGTERM",
|
||||
limit=1,
|
||||
post_kill_wait_seconds=2.0,
|
||||
),
|
||||
id="runtime_process_chain_sigterm",
|
||||
),
|
||||
pytest.param(
|
||||
make_process_kill_fault_injector(
|
||||
grep_patterns=RUNTIME_WORKER_PATTERN,
|
||||
signal_name="SIGKILL",
|
||||
limit=1,
|
||||
post_kill_wait_seconds=2.0,
|
||||
),
|
||||
id="runtime_process_chain_sigkill",
|
||||
),
|
||||
]
|
||||
|
||||
VOXCPM2_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLOY_CONFIGS_DIR)
|
||||
INFLIGHT_INJECTION_REQUEST_RATE = 0.3
|
||||
INFLIGHT_INJECTION_REQUEST_COUNT = 10
|
||||
_SPEECH_FF_HTTP_TIMEOUT_SEC = 30
|
||||
_SPEECH_FF_MAX_ELAPSED_SEC = 30
|
||||
|
||||
|
||||
def _speech_request_config(omni_server: Any) -> dict[str, Any]:
|
||||
return {
|
||||
"model": omni_server.model,
|
||||
"input": "The weather is nice today, perfect for a walk in the park.",
|
||||
"stream": False,
|
||||
"timeout": 120.0,
|
||||
"response_format": "wav",
|
||||
"voice": "default",
|
||||
}
|
||||
|
||||
|
||||
def _submit_speech_request(openai_client: OpenAIClientHandler, omni_server: Any) -> None:
|
||||
openai_client.send_audio_speech_request(_speech_request_config(omni_server), request_num=1)
|
||||
|
||||
|
||||
def _oom_speech_request_config(omni_server: Any) -> dict[str, Any]:
|
||||
base = "The weather is nice today, perfect for a walk in the park. "
|
||||
return {
|
||||
"model": omni_server.model,
|
||||
"input": base * 80,
|
||||
"max_new_tokens": 4096,
|
||||
"ref_audio": load_test_audio_data_url("qwen3_tts/clone_2.wav"),
|
||||
"stream": False,
|
||||
"timeout": 300.0,
|
||||
"response_format": "wav",
|
||||
"voice": "default",
|
||||
}
|
||||
|
||||
|
||||
def _submit_oom_speech_request(openai_client: OpenAIClientHandler, omni_server: Any) -> None:
|
||||
openai_client.send_audio_speech_request(_oom_speech_request_config(omni_server), request_num=1)
|
||||
|
||||
|
||||
def _assert_post_fault_speech_fast_fail(host: str, port: int, *, model: str, scenario: str) -> None:
|
||||
payload = {
|
||||
"model": model,
|
||||
"input": "post-fault fast-fail check",
|
||||
"voice": "default",
|
||||
}
|
||||
start = time.monotonic()
|
||||
try:
|
||||
status, body = post_json_raw(host, port, "/v1/audio/speech", payload, timeout_sec=20)
|
||||
elapsed = time.monotonic() - start
|
||||
assert elapsed < 15, f"[{scenario} fast_fail] /v1/audio/speech did not fail fast: {elapsed:.2f}s"
|
||||
assert status >= 500, (
|
||||
f"[{scenario} fast_fail] expected server-side error after fault, got status={status}, body={body[:200]!r}"
|
||||
)
|
||||
except Exception: # noqa: BLE001
|
||||
elapsed = time.monotonic() - start
|
||||
assert elapsed < 15, f"[{scenario} fast_fail] exception was too slow after fault: {elapsed:.2f}s"
|
||||
|
||||
|
||||
def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -> None:
|
||||
deadline = time.monotonic() + 20.0
|
||||
last_observation = ""
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
status, body = get_health_raw(host, port, timeout_sec=5)
|
||||
last_observation = f"http={status}, body={body[:200]!r}"
|
||||
if status == 503:
|
||||
return
|
||||
except Exception as exc: # noqa: BLE001
|
||||
last_observation = f"exception={exc!r}"
|
||||
return
|
||||
time.sleep(0.5)
|
||||
pytest.fail(f"[{scenario} health] no terminal post-fault health observed: {last_observation}")
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@pytest.mark.skip(reason="issue#4285")
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
|
||||
def test_reliability_fault_gpu_oom_speech_request_failure(
|
||||
omni_server_function,
|
||||
openai_client_function,
|
||||
) -> None:
|
||||
stage_config_path = getattr(omni_server_function, "stage_config_path", None)
|
||||
device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
|
||||
handle = inject_gpu_oom(
|
||||
device=device_spec,
|
||||
target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
|
||||
hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
|
||||
startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
|
||||
strict=OOM_INJECTION_CONFIG["strict"],
|
||||
)
|
||||
try:
|
||||
try:
|
||||
_submit_oom_speech_request(openai_client_function, omni_server_function)
|
||||
except Exception as exc:
|
||||
assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
|
||||
else:
|
||||
pytest.fail("expected /v1/audio/speech request failure during GPU OOM injection")
|
||||
finally:
|
||||
stop_gpu_oom_hogs(handle)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@pytest.mark.skip(reason="issue#4285")
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
|
||||
def test_reliability_fault_gpu_oom_speech_error_contract(
|
||||
omni_server_function,
|
||||
) -> None:
|
||||
stage_config_path = getattr(omni_server_function, "stage_config_path", None)
|
||||
device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
|
||||
handle = inject_gpu_oom(
|
||||
device=device_spec,
|
||||
target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
|
||||
hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
|
||||
startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
|
||||
strict=OOM_INJECTION_CONFIG["strict"],
|
||||
)
|
||||
host = omni_server_function.host
|
||||
port = omni_server_function.port
|
||||
oom_config = _oom_speech_request_config(omni_server_function)
|
||||
payload = {
|
||||
"model": oom_config["model"],
|
||||
"input": oom_config["input"],
|
||||
"voice": oom_config["voice"],
|
||||
"max_new_tokens": oom_config["max_new_tokens"],
|
||||
"ref_audio": oom_config["ref_audio"],
|
||||
}
|
||||
try:
|
||||
status, body = post_json_raw(host, port, "/v1/audio/speech", payload, timeout_sec=300)
|
||||
finally:
|
||||
stop_gpu_oom_hogs(handle)
|
||||
|
||||
assert status >= 500, f"expected speech runtime error under OOM, got status={status}"
|
||||
speech_error = extract_openai_error_contract_from_bytes(body)
|
||||
assert speech_error is not None, f"speech error payload not OpenAI-compatible: {body[:300]!r}"
|
||||
assert "code" in speech_error, f"speech error lacks code field: {speech_error!r}"
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
|
||||
@pytest.mark.parametrize(
|
||||
"fault_injector",
|
||||
WORKER_SIGNAL_FAULT_PARAMS,
|
||||
indirect=True,
|
||||
)
|
||||
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_request_failure(
|
||||
omni_server_after_fault_function,
|
||||
openai_client_function,
|
||||
) -> None:
|
||||
try:
|
||||
_submit_speech_request(openai_client_function, omni_server_after_fault_function)
|
||||
except Exception as exc:
|
||||
assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
|
||||
else:
|
||||
pytest.fail("expected /v1/audio/speech request failure after process-kill injection")
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
|
||||
@pytest.mark.parametrize(
|
||||
"fault_injector",
|
||||
WORKER_SIGNAL_FAULT_PARAMS,
|
||||
indirect=True,
|
||||
)
|
||||
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_fast_fail_and_concurrent(
|
||||
omni_server_after_fault_function,
|
||||
openai_client_function,
|
||||
) -> None:
|
||||
host = omni_server_after_fault_function.host
|
||||
port = omni_server_after_fault_function.port
|
||||
model = omni_server_after_fault_function.model
|
||||
payload = {
|
||||
"model": model,
|
||||
"input": "Say hello in one short sentence.",
|
||||
"voice": "default",
|
||||
}
|
||||
|
||||
start = time.monotonic()
|
||||
try:
|
||||
status, body = post_json_raw(
|
||||
host,
|
||||
port,
|
||||
"/v1/audio/speech",
|
||||
payload,
|
||||
timeout_sec=_SPEECH_FF_HTTP_TIMEOUT_SEC,
|
||||
)
|
||||
elapsed = time.monotonic() - start
|
||||
assert elapsed <= _SPEECH_FF_MAX_ELAPSED_SEC + 1.0, (
|
||||
f"[process_kill fast_fail] /v1/audio/speech did not fail fast after fault: {elapsed:.2f}s"
|
||||
)
|
||||
assert status >= 500, (
|
||||
"[process_kill fast_fail] expected server-side failure after fault, "
|
||||
f"got status={status}, body={body[:200]!r}"
|
||||
)
|
||||
except Exception:
|
||||
elapsed = time.monotonic() - start
|
||||
assert elapsed <= _SPEECH_FF_MAX_ELAPSED_SEC + 1.0, (
|
||||
f"[process_kill fast_fail] /v1/audio/speech exception was too slow after fault: {elapsed:.2f}s"
|
||||
)
|
||||
|
||||
start = time.monotonic()
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
|
||||
futures = [
|
||||
executor.submit(_submit_speech_request, openai_client_function, omni_server_after_fault_function)
|
||||
for _ in range(3)
|
||||
]
|
||||
done, pending = concurrent.futures.wait(
|
||||
futures,
|
||||
timeout=30,
|
||||
return_when=concurrent.futures.ALL_COMPLETED,
|
||||
)
|
||||
|
||||
elapsed = time.monotonic() - start
|
||||
assert not pending, f"[process_kill concurrent] some fault-time speech requests hung: pending={len(pending)}"
|
||||
assert elapsed < 30, f"[process_kill concurrent] fault-time speech request convergence is too slow: {elapsed:.2f}s"
|
||||
|
||||
fault_observed = False
|
||||
for future in done:
|
||||
try:
|
||||
future.result()
|
||||
except Exception as exc:
|
||||
fault_observed = True
|
||||
assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
|
||||
assert fault_observed, (
|
||||
"[process_kill concurrent] expected at least one /v1/audio/speech request to fail after fault"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
|
||||
@pytest.mark.parametrize(
|
||||
"fault_injector",
|
||||
WORKER_SIGNAL_FAULT_PARAMS,
|
||||
indirect=True,
|
||||
)
|
||||
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_worker_with_load_request_failure(
|
||||
omni_server_function,
|
||||
openai_client_function,
|
||||
fault_injector: FaultInjector,
|
||||
) -> None:
|
||||
scenario = "kill_worker_with_load"
|
||||
load_result = run_fault_injection_with_rate_load(
|
||||
submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
|
||||
inject_fault=lambda: fault_injector(omni_server_function),
|
||||
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
|
||||
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
|
||||
completion_timeout_sec=120.0,
|
||||
)
|
||||
assert load_result["failure_observed"], (
|
||||
f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
|
||||
)
|
||||
host = omni_server_function.host
|
||||
port = omni_server_function.port
|
||||
_assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
|
||||
_assert_post_fault_health_terminal(host, port, scenario=scenario)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
|
||||
@pytest.mark.parametrize("signal_name", SERVE_SIGNAL_PARAMS)
|
||||
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_serve_root_no_load_fast_fail_and_cleanup(
|
||||
omni_server_function,
|
||||
signal_name: str,
|
||||
) -> None:
|
||||
scenario = f"kill_serve_root_no_load_{signal_name.lower()}"
|
||||
injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
|
||||
injector(omni_server_function)
|
||||
host = omni_server_function.host
|
||||
port = omni_server_function.port
|
||||
_assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
|
||||
_assert_post_fault_health_terminal(host, port, scenario=scenario)
|
||||
assert_no_server_tree_process_residual_and_gpu_release(
|
||||
omni_server_function,
|
||||
scenario=scenario,
|
||||
timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
|
||||
@pytest.mark.parametrize("signal_name", SERVE_WITH_LOAD_SIGNAL_PARAMS)
|
||||
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup(
|
||||
omni_server_function,
|
||||
openai_client_function,
|
||||
signal_name: str,
|
||||
) -> None:
|
||||
scenario = f"kill_serve_root_with_load_{signal_name.lower()}"
|
||||
injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
|
||||
load_result = run_fault_injection_with_rate_load(
|
||||
submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
|
||||
inject_fault=lambda: injector(omni_server_function),
|
||||
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
|
||||
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
|
||||
completion_timeout_sec=300.0,
|
||||
)
|
||||
assert load_result["failure_observed"], (
|
||||
f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
|
||||
)
|
||||
host = omni_server_function.host
|
||||
port = omni_server_function.port
|
||||
_assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
|
||||
_assert_post_fault_health_terminal(host, port, scenario=scenario)
|
||||
assert_no_server_tree_process_residual_and_gpu_release(
|
||||
omni_server_function,
|
||||
scenario=scenario,
|
||||
timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
|
||||
@pytest.mark.parametrize("signal_name", TREE_SIGNAL_PARAMS)
|
||||
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup(
|
||||
omni_server_function,
|
||||
signal_name: str,
|
||||
) -> None:
|
||||
scenario = f"kill_serve_tree_no_load_{signal_name.lower()}"
|
||||
injector = make_server_tree_kill_fault_injector(
|
||||
signal_name=signal_name,
|
||||
post_kill_wait_seconds=2.0,
|
||||
inter_kill_wait_seconds=0.1,
|
||||
)
|
||||
injector(omni_server_function)
|
||||
host = omni_server_function.host
|
||||
port = omni_server_function.port
|
||||
_assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
|
||||
_assert_post_fault_health_terminal(host, port, scenario=scenario)
|
||||
assert_no_server_tree_process_residual_and_gpu_release(
|
||||
omni_server_function,
|
||||
scenario=scenario,
|
||||
timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
|
||||
@pytest.mark.parametrize("signal_name", TREE_WITH_LOAD_SIGNAL_PARAMS)
|
||||
@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
|
||||
omni_server_function,
|
||||
openai_client_function,
|
||||
signal_name: str,
|
||||
) -> None:
|
||||
scenario = f"kill_serve_tree_with_load_{signal_name.lower()}"
|
||||
injector = make_server_tree_kill_fault_injector(
|
||||
signal_name=signal_name,
|
||||
post_kill_wait_seconds=2.0,
|
||||
inter_kill_wait_seconds=0.1,
|
||||
)
|
||||
load_result = run_fault_injection_with_rate_load(
|
||||
submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function),
|
||||
inject_fault=lambda: injector(omni_server_function),
|
||||
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
|
||||
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
|
||||
completion_timeout_sec=120.0,
|
||||
)
|
||||
assert load_result["failure_observed"], (
|
||||
f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}"
|
||||
)
|
||||
host = omni_server_function.host
|
||||
port = omni_server_function.port
|
||||
_assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario)
|
||||
_assert_post_fault_health_terminal(host, port, scenario=scenario)
|
||||
assert_no_server_tree_process_residual_and_gpu_release(
|
||||
omni_server_function,
|
||||
scenario=scenario,
|
||||
timeout_sec=worker_residual_timeout_after_kill_signal(signal_name),
|
||||
)
|
||||
@@ -1,3 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""
|
||||
Wan2.2 reliability integration tests.
|
||||
"""
|
||||
@@ -183,7 +186,7 @@ def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -
|
||||
@pytest.mark.slow
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
|
||||
def test_reliability_fault_gpu_oom_video_large_request_failure(omni_server_function, openai_client_function) -> None:
|
||||
def test_reliability_fault_gpu_oom_video_large_request_failure(omni_server_function, online_client_function) -> None:
|
||||
stage_config_path = getattr(omni_server_function, "stage_config_path", None)
|
||||
device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
|
||||
handle = inject_gpu_oom(
|
||||
@@ -211,7 +214,7 @@ def test_reliability_fault_gpu_oom_video_large_request_failure(omni_server_funct
|
||||
"stream": False,
|
||||
}
|
||||
try:
|
||||
openai_client_function.send_video_diffusion_request(request_config, request_num=1)
|
||||
online_client_function.send_video_diffusion_request(request_config, request_num=1)
|
||||
except Exception as exc:
|
||||
assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
|
||||
else:
|
||||
@@ -230,11 +233,11 @@ def test_reliability_fault_gpu_oom_video_large_request_failure(omni_server_funct
|
||||
@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_request_failure(
|
||||
omni_server_after_fault_function,
|
||||
openai_client_function,
|
||||
online_client_function,
|
||||
) -> None:
|
||||
request_config = _video_request_config()
|
||||
try:
|
||||
openai_client_function.send_video_diffusion_request(request_config, request_num=1)
|
||||
online_client_function.send_video_diffusion_request(request_config, request_num=1)
|
||||
except Exception as exc:
|
||||
assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
|
||||
else:
|
||||
@@ -251,7 +254,7 @@ def test_reliability_fault_process_kill_request_failure(
|
||||
@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
|
||||
omni_server_after_fault_function,
|
||||
openai_client_function,
|
||||
online_client_function,
|
||||
) -> None:
|
||||
"""Black-box: after process kill, /v1/videos fails fast and concurrent calls don't hang.
|
||||
|
||||
@@ -308,7 +311,7 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
|
||||
start = time.monotonic()
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
|
||||
futures = [
|
||||
executor.submit(openai_client_function.send_video_diffusion_request, request_config, 1) for _ in range(3)
|
||||
executor.submit(online_client_function.send_video_diffusion_request, request_config, 1) for _ in range(3)
|
||||
]
|
||||
done, pending = concurrent.futures.wait(
|
||||
futures,
|
||||
@@ -367,13 +370,13 @@ def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
|
||||
@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_worker_with_load_request_failure(
|
||||
omni_server_function,
|
||||
openai_client_function,
|
||||
online_client_function,
|
||||
fault_injector: FaultInjector,
|
||||
) -> None:
|
||||
request_config = _video_request_config()
|
||||
scenario = "kill_worker_with_load"
|
||||
load_result = run_fault_injection_with_rate_load(
|
||||
submit_request=lambda: openai_client_function.send_video_diffusion_request(request_config, request_num=1),
|
||||
submit_request=lambda: online_client_function.send_video_diffusion_request(request_config, request_num=1),
|
||||
inject_fault=lambda: fault_injector(omni_server_function),
|
||||
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
|
||||
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
|
||||
@@ -516,7 +519,7 @@ def test_reliability_video_oom_recovers_after_fault_removed(
|
||||
@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup(
|
||||
omni_server_function,
|
||||
openai_client_function,
|
||||
online_client_function,
|
||||
signal_name: str,
|
||||
) -> None:
|
||||
"""Black-box: during in-flight video requests, kill serve root and verify post-fault behavior/cleanup."""
|
||||
@@ -524,7 +527,7 @@ def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_clean
|
||||
scenario = f"kill_serve_root_with_load_{signal_name.lower()}"
|
||||
injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0)
|
||||
load_result = run_fault_injection_with_rate_load(
|
||||
submit_request=lambda: openai_client_function.send_video_diffusion_request(request_config, request_num=1),
|
||||
submit_request=lambda: online_client_function.send_video_diffusion_request(request_config, request_num=1),
|
||||
inject_fault=lambda: injector(omni_server_function),
|
||||
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
|
||||
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
|
||||
@@ -601,7 +604,7 @@ def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup(
|
||||
@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
|
||||
def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
|
||||
omni_server_function,
|
||||
openai_client_function,
|
||||
online_client_function,
|
||||
signal_name: str,
|
||||
) -> None:
|
||||
"""Black-box: during in-flight video requests, kill server process tree and verify post-fault behavior/cleanup."""
|
||||
@@ -613,7 +616,7 @@ def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
|
||||
inter_kill_wait_seconds=0.1,
|
||||
)
|
||||
load_result = run_fault_injection_with_rate_load(
|
||||
submit_request=lambda: openai_client_function.send_video_diffusion_request(request_config, request_num=1),
|
||||
submit_request=lambda: online_client_function.send_video_diffusion_request(request_config, request_num=1),
|
||||
inject_fault=lambda: injector(omni_server_function),
|
||||
num_requests=INFLIGHT_INJECTION_REQUEST_COUNT,
|
||||
request_rate=INFLIGHT_INJECTION_REQUEST_RATE,
|
||||
|
||||
@@ -37,7 +37,7 @@ def omni_server(request: pytest.FixtureRequest):
|
||||
with _omni_server_lock:
|
||||
# Same tuple and CLI composition as ``tests/dfx/perf/scripts/run_benchmark.py``;
|
||||
# ``serve_args`` from JSON are folded into ``extra_cli_args`` inside
|
||||
# ``create_unique_server_params``.
|
||||
# ``create_unique_server_pytest_params`` (JSON ``mark`` → pytest marks).
|
||||
test_name, model, deploy_path, stage_overrides, extra_cli_args, use_omni = request.param
|
||||
|
||||
print(f"Starting OmniServer with test: {test_name}, model: {model}")
|
||||
|
||||
@@ -13,7 +13,7 @@ import pytest
|
||||
from tests.dfx.conftest import (
|
||||
create_benchmark_indices,
|
||||
create_test_parameter_mapping,
|
||||
create_unique_server_params,
|
||||
create_unique_server_pytest_params,
|
||||
load_configs,
|
||||
)
|
||||
from tests.dfx.stability.helpers import _run_one_diffusion_batch, run_stability_benchmark_loop
|
||||
@@ -29,7 +29,7 @@ try:
|
||||
except FileNotFoundError:
|
||||
BENCHMARK_CONFIGS = []
|
||||
|
||||
test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
|
||||
test_params = create_unique_server_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
|
||||
server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {}
|
||||
benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
|
||||
|
||||
|
||||
@@ -13,7 +13,7 @@ import pytest
|
||||
from tests.dfx.conftest import (
|
||||
create_benchmark_indices,
|
||||
create_test_parameter_mapping,
|
||||
create_unique_server_params,
|
||||
create_unique_server_pytest_params,
|
||||
load_configs,
|
||||
)
|
||||
from tests.dfx.stability.helpers import _run_one_vllm_bench_batch, run_stability_benchmark_loop
|
||||
@@ -29,7 +29,7 @@ try:
|
||||
except FileNotFoundError:
|
||||
BENCHMARK_CONFIGS = []
|
||||
|
||||
test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
|
||||
test_params = create_unique_server_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
|
||||
server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {}
|
||||
benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
|
||||
|
||||
|
||||
@@ -13,7 +13,7 @@ import pytest
|
||||
from tests.dfx.conftest import (
|
||||
create_benchmark_indices,
|
||||
create_test_parameter_mapping,
|
||||
create_unique_server_params,
|
||||
create_unique_server_pytest_params,
|
||||
load_configs,
|
||||
)
|
||||
from tests.dfx.stability.helpers import _run_one_vllm_bench_batch, run_stability_benchmark_loop
|
||||
@@ -29,7 +29,7 @@ try:
|
||||
except FileNotFoundError:
|
||||
BENCHMARK_CONFIGS = []
|
||||
|
||||
test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
|
||||
test_params = create_unique_server_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
|
||||
server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {}
|
||||
benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
|
||||
|
||||
|
||||
@@ -13,7 +13,7 @@ import pytest
|
||||
from tests.dfx.conftest import (
|
||||
create_benchmark_indices,
|
||||
create_test_parameter_mapping,
|
||||
create_unique_server_params,
|
||||
create_unique_server_pytest_params,
|
||||
load_configs,
|
||||
)
|
||||
from tests.dfx.stability.helpers import _run_one_diffusion_batch, run_stability_benchmark_loop
|
||||
@@ -29,7 +29,7 @@ try:
|
||||
except FileNotFoundError:
|
||||
BENCHMARK_CONFIGS = []
|
||||
|
||||
test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
|
||||
test_params = create_unique_server_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
|
||||
server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {}
|
||||
benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
|
||||
|
||||
|
||||
@@ -1,69 +0,0 @@
|
||||
"""
|
||||
VoxCPM2 stability: OmniServer + ``vllm bench serve --omni`` for a fixed duration.
|
||||
|
||||
Configuration: ``tests/dfx/stability/tests/test_voxcpm2.json``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.dfx.conftest import (
|
||||
create_benchmark_indices,
|
||||
create_test_parameter_mapping,
|
||||
create_unique_server_params,
|
||||
load_configs,
|
||||
)
|
||||
from tests.dfx.stability.helpers import _run_one_vllm_bench_batch, run_stability_benchmark_loop
|
||||
|
||||
STABILITY_DIR = Path(__file__).resolve().parent.parent
|
||||
DEPLOY_CONFIGS_DIR = STABILITY_DIR / "deploy"
|
||||
CONFIG_FILE_PATH = str(STABILITY_DIR / "tests" / "test_voxcpm2.json")
|
||||
DEFAULT_NUM_PROMPTS_PER_BATCH = 20
|
||||
STABILITY_SERVER_TIMEOUT_ARGS = ["--stage-init-timeout", "600"]
|
||||
|
||||
try:
|
||||
BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH)
|
||||
except FileNotFoundError:
|
||||
BENCHMARK_CONFIGS = []
|
||||
|
||||
test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
|
||||
server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {}
|
||||
benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.tts
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
@pytest.mark.parametrize("stability_benchmark_params", benchmark_indices, indirect=True)
|
||||
def test_stability_voxcpm2(omni_server, stability_benchmark_params):
|
||||
test_name = stability_benchmark_params["test_name"]
|
||||
params = stability_benchmark_params["params"]
|
||||
duration_sec = params.get("duration_sec", 300)
|
||||
num_prompts_per_batch = params.get("num_prompts_per_batch", DEFAULT_NUM_PROMPTS_PER_BATCH)
|
||||
request_rate = params.get("request_rate")
|
||||
max_concurrency = params.get("max_concurrency")
|
||||
|
||||
bench_params = {
|
||||
k: v
|
||||
for k, v in params.items()
|
||||
if k not in ("duration_sec", "request_rate", "max_concurrency", "num_prompts_per_batch")
|
||||
}
|
||||
|
||||
result = run_stability_benchmark_loop(
|
||||
host=omni_server.host,
|
||||
port=omni_server.port,
|
||||
model=omni_server.model,
|
||||
duration_sec=duration_sec,
|
||||
params=bench_params,
|
||||
request_rate=request_rate,
|
||||
max_concurrency=max_concurrency,
|
||||
result_dir=str(STABILITY_DIR),
|
||||
num_prompts_per_batch=num_prompts_per_batch,
|
||||
run_one_batch=_run_one_vllm_bench_batch,
|
||||
)
|
||||
|
||||
assert result.get("failed", 0) == 0, f"[{test_name}] Failed requests detected: {result.get('errors', [])}"
|
||||
assert result.get("completed", 0) > 0, f"[{test_name}] No requests completed"
|
||||
@@ -13,7 +13,7 @@ import pytest
|
||||
from tests.dfx.conftest import (
|
||||
create_benchmark_indices,
|
||||
create_test_parameter_mapping,
|
||||
create_unique_server_params,
|
||||
create_unique_server_pytest_params,
|
||||
load_configs,
|
||||
)
|
||||
from tests.dfx.stability.helpers import _run_one_diffusion_batch, run_stability_benchmark_loop
|
||||
@@ -29,7 +29,7 @@ try:
|
||||
except FileNotFoundError:
|
||||
BENCHMARK_CONFIGS = []
|
||||
|
||||
test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
|
||||
test_params = create_unique_server_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else []
|
||||
server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {}
|
||||
benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping)
|
||||
|
||||
|
||||
@@ -1,6 +1,18 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_hunyuan_image_stability_tp2_fp8_sp2",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H800"
|
||||
},
|
||||
"num_cards": 2
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"diffusion"
|
||||
],
|
||||
"server_params": {
|
||||
"model": "tencent/HunyuanImage-3.0-Instruct",
|
||||
"serve_args": {
|
||||
|
||||
@@ -1,6 +1,18 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_qwen3_omni_stability_async_chunk",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H800"
|
||||
},
|
||||
"num_cards": 2
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"omni"
|
||||
],
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
||||
"extra_cli_args": ["--async-chunk"]
|
||||
|
||||
@@ -1,6 +1,19 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_qwen3_tts_stability",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H800",
|
||||
"npu": "A3"
|
||||
},
|
||||
"num_cards": 2
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"tts"
|
||||
],
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
|
||||
},
|
||||
|
||||
@@ -1,6 +1,18 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_qwen_image_stability",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H800"
|
||||
},
|
||||
"num_cards": 1
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"diffusion"
|
||||
],
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen-Image"
|
||||
},
|
||||
|
||||
@@ -1,35 +0,0 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_voxcpm2_stability",
|
||||
"server_params": {
|
||||
"model": "openbmb/VoxCPM2",
|
||||
"trust_remote_code": true
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"dataset_name": "random",
|
||||
"backend": "openai-audio-speech",
|
||||
"endpoint": "/v1/audio/speech",
|
||||
"trust_remote_code": true,
|
||||
"duration_sec": 86400,
|
||||
"max_concurrency": 2,
|
||||
"num_prompts_per_batch": 50,
|
||||
"random_input_len": {
|
||||
"min": 0,
|
||||
"max": 1000
|
||||
},
|
||||
"random_output_len": {
|
||||
"min": 0,
|
||||
"max": 1000
|
||||
},
|
||||
"random_range_ratio": 0.0,
|
||||
"extra_body": {
|
||||
"voice": "default",
|
||||
"language": "English"
|
||||
},
|
||||
"ignore_eos": true,
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -1,6 +1,19 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_wan22_i2v_stability_v1_videos",
|
||||
"test_name": "test_wan22_i2v_stability_ulysses2_vae_patch2_hsdp_slicing_tiling",
|
||||
"description": "Ulysses=2 + VAE patch parallel=2 + HSDP + VAE slicing + tiling",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"cuda": "H800"
|
||||
},
|
||||
"num_cards": 2
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"diffusion"
|
||||
],
|
||||
"server_params": {
|
||||
"model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
|
||||
"serve_args": {
|
||||
@@ -27,5 +40,45 @@
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_wan22_i2v_stability_usp2_vae_patch2_hsdp_slicing",
|
||||
"description": "USP=2 + VAE patch parallel=2 + HSDP + VAE slicing",
|
||||
"mark": [
|
||||
{
|
||||
"hardware_marks": {
|
||||
"res": {
|
||||
"npu": "A3"
|
||||
},
|
||||
"num_cards": 2
|
||||
}
|
||||
},
|
||||
"local_model",
|
||||
"diffusion"
|
||||
],
|
||||
"server_params": {
|
||||
"model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
|
||||
"serve_args": {
|
||||
"usp": 2,
|
||||
"vae-patch-parallel-size": 2,
|
||||
"use-hsdp": true,
|
||||
"vae-use-slicing": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"dataset": "random",
|
||||
"task": "i2v",
|
||||
"backend": "v1/videos",
|
||||
"duration_sec": 86400,
|
||||
"max_concurrency": 1,
|
||||
"num_prompts_per_batch": 50,
|
||||
"enable_negative_prompt": true,
|
||||
"random_request_config": [
|
||||
{"width": 832, "height": 480, "num_inference_steps": 4, "num_frames": 81, "fps": 16, "weight": 0.5},
|
||||
{"width": 1280, "height": 720, "num_inference_steps": 4, "num_frames": 121, "fps": 16, "weight": 0.5}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -18,6 +18,7 @@ from vllm_omni.experimental.ar_diffusion.capability import (
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVConfig
|
||||
from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
|
||||
from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionTickRequest
|
||||
|
||||
BLOCK = 16
|
||||
POS = "positive"
|
||||
@@ -56,6 +57,21 @@ def dreamzero_like_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
|
||||
)
|
||||
|
||||
|
||||
def tiny_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
|
||||
return ARDiffusionKVCacheSpec(
|
||||
num_layers=1,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
tokens_per_frame=1,
|
||||
frames_per_block=3,
|
||||
window_frames=3,
|
||||
sink_frames=3,
|
||||
kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
|
||||
session_capacity=capacity,
|
||||
cross_attention=(ARDiffusionCrossAttentionKVSpec("text", 2),),
|
||||
)
|
||||
|
||||
|
||||
class CapablePipeline:
|
||||
def __init__(self, spec: ARDiffusionKVCacheSpec) -> None:
|
||||
self.spec = spec
|
||||
@@ -103,6 +119,7 @@ def make_runner(
|
||||
*,
|
||||
available_bytes: int = 1 << 28,
|
||||
step_execution: bool = False,
|
||||
gpu_memory_fraction: float = 0.1,
|
||||
) -> ARDiffusionModelRunner:
|
||||
runner = object.__new__(ARDiffusionModelRunner)
|
||||
runner.od_config = SimpleNamespace(
|
||||
@@ -113,7 +130,10 @@ def make_runner(
|
||||
)
|
||||
runner.device = torch.device("cpu")
|
||||
runner.pipeline = pipeline
|
||||
runner.ar_diffusion_kv_config = ARDiffusionKVConfig(enable=True)
|
||||
runner.ar_diffusion_kv_config = ARDiffusionKVConfig(
|
||||
enable=True,
|
||||
gpu_memory_fraction=gpu_memory_fraction,
|
||||
)
|
||||
runner.kv_cache = None
|
||||
runner._ar_diffusion_capability = None
|
||||
runner._ar_diffusion_kv_cache_spec = None
|
||||
@@ -132,6 +152,16 @@ def commit_one_frame(runner: ARDiffusionModelRunner, session_id: str, kv_branch:
|
||||
return state
|
||||
|
||||
|
||||
def commit_one_block(runner: ARDiffusionModelRunner, session_id: str, kv_branch: str):
|
||||
state = runner._get_or_create_session(session_id)
|
||||
assert runner.kv_cache is not None
|
||||
seq_len = BLOCK * runner.kv_cache.frames_per_block
|
||||
ctx = state.get_kv_caches(kv_branch, seq_len=seq_len, commit_current=True)[0].forward_ctx
|
||||
ctx.ensure_video_slots(torch.device("cpu"))
|
||||
state.commit_paged_context(kv_branch)
|
||||
return state
|
||||
|
||||
|
||||
def test_ar_runner_rejects_pipeline_without_capability():
|
||||
runner = object.__new__(ARDiffusionModelRunner)
|
||||
runner.od_config = SimpleNamespace(max_num_seqs=1)
|
||||
@@ -140,6 +170,43 @@ def test_ar_runner_rejects_pipeline_without_capability():
|
||||
runner._preallocate_kv_cache(available_bytes=1 << 20)
|
||||
|
||||
|
||||
def test_runner_uses_typed_tick_as_authoritative_session_contract():
|
||||
tick = ARDiffusionTickRequest(
|
||||
session_id="world-7",
|
||||
request_id="request-3",
|
||||
chunk_index=3,
|
||||
reset=True,
|
||||
)
|
||||
req = SimpleNamespace(
|
||||
request_id="request-3",
|
||||
sampling_params=SimpleNamespace(extra_args=tick.to_extra_args()),
|
||||
)
|
||||
|
||||
session_id, extra_args, parsed = ARDiffusionModelRunner._request_session(req)
|
||||
|
||||
assert session_id == "world-7"
|
||||
assert extra_args == tick.to_extra_args()
|
||||
assert parsed == tick
|
||||
|
||||
|
||||
def test_runner_keeps_engine_request_id_separate_from_tick_correlation_id():
|
||||
tick = ARDiffusionTickRequest(
|
||||
session_id="world-7",
|
||||
request_id="client-request-3",
|
||||
chunk_index=3,
|
||||
)
|
||||
req = SimpleNamespace(
|
||||
request_id="engine-request-uuid",
|
||||
sampling_params=SimpleNamespace(extra_args=tick.to_extra_args()),
|
||||
)
|
||||
|
||||
session_id, _, parsed = ARDiffusionModelRunner._request_session(req)
|
||||
|
||||
assert session_id == "world-7"
|
||||
assert req.request_id == "engine-request-uuid"
|
||||
assert parsed.request_id == "client-request-3"
|
||||
|
||||
|
||||
def test_lingbot_like_single_branch_session_reuse_reset_and_close():
|
||||
pipeline = CapablePipeline(lingbot_like_spec())
|
||||
runner = make_runner(pipeline)
|
||||
@@ -170,6 +237,31 @@ def test_lingbot_like_single_branch_session_reuse_reset_and_close():
|
||||
assert pipeline.closes == ["s1"]
|
||||
|
||||
|
||||
def test_lingbot_like_interleaved_sessions_keep_independent_kv_partitions():
|
||||
runner = make_runner(CapablePipeline(lingbot_like_spec(capacity=2)))
|
||||
kv = runner.kv_cache
|
||||
assert kv is not None
|
||||
|
||||
session_a = commit_one_block(runner, "world-a", "main")
|
||||
a0_blocks = kv.window_block_ids(session_a.adapter("main"))
|
||||
session_b = commit_one_block(runner, "world-b", "main")
|
||||
b0_blocks = kv.window_block_ids(session_b.adapter("main"))
|
||||
session_a_again = commit_one_block(runner, "world-a", "main")
|
||||
a1_blocks = kv.window_block_ids(session_a_again.adapter("main"))
|
||||
|
||||
assert session_a_again is session_a
|
||||
assert session_a.adapter("main").request_id == "ar::world-a::main"
|
||||
assert session_b.adapter("main").request_id == "ar::world-b::main"
|
||||
assert session_a.adapter("main").completed_chunks == 6
|
||||
assert session_b.adapter("main").completed_chunks == 3
|
||||
assert len(a0_blocks) == 3
|
||||
assert len(b0_blocks) == 3
|
||||
assert len(a1_blocks) == 6
|
||||
assert set(a0_blocks) <= set(a1_blocks)
|
||||
assert set(a1_blocks).isdisjoint(b0_blocks)
|
||||
assert tuple(runner._sessions) == ("world-b", "world-a")
|
||||
|
||||
|
||||
def test_lingbot_like_sink_survives_sliding_window_eviction():
|
||||
runner = make_runner(CapablePipeline(lingbot_like_spec()))
|
||||
kv = runner.kv_cache
|
||||
@@ -198,12 +290,12 @@ def test_dreamzero_like_two_branches_are_independent():
|
||||
|
||||
|
||||
def test_lru_eviction_releases_blocks_and_notifies_pipeline():
|
||||
pipeline = CapablePipeline(lingbot_like_spec(capacity=8))
|
||||
pipeline = CapablePipeline(lingbot_like_spec(capacity=2))
|
||||
runner = make_runner(pipeline)
|
||||
kv = runner.kv_cache
|
||||
assert kv is not None
|
||||
assert runner._session_capacity == 1
|
||||
assert kv.session_capacity == 1
|
||||
assert runner._session_capacity == 2
|
||||
assert kv.session_capacity == 2
|
||||
free_total = kv.manager.block_pool.get_num_free_blocks()
|
||||
old = commit_one_frame(runner, "old", "main")
|
||||
k = torch.randn(1, 8, 4, 64)
|
||||
@@ -211,13 +303,56 @@ def test_lru_eviction_releases_blocks_and_notifies_pipeline():
|
||||
assert kv.manager.block_pool.get_num_free_blocks() < free_total
|
||||
|
||||
runner._get_or_create_session("new")
|
||||
assert tuple(runner._sessions) == ("old", "new")
|
||||
runner._get_or_create_session("newest")
|
||||
|
||||
assert tuple(runner._sessions) == ("new",)
|
||||
assert tuple(runner._sessions) == ("new", "newest")
|
||||
assert pipeline.closes == ["old"]
|
||||
assert "old" not in kv._cross_sessions
|
||||
assert kv.manager.block_pool.get_num_free_blocks() == free_total
|
||||
|
||||
|
||||
def test_budget_reduced_capacity_drives_runner_lru():
|
||||
pipeline = CapablePipeline(tiny_spec(capacity=2))
|
||||
# One tiny LingBot-like session requires 128 bytes; two require 192.
|
||||
runner = make_runner(
|
||||
pipeline,
|
||||
available_bytes=128,
|
||||
gpu_memory_fraction=1.0,
|
||||
)
|
||||
kv = runner.kv_cache
|
||||
assert kv is not None
|
||||
assert kv.requested_session_capacity == 2
|
||||
assert kv.session_capacity == 1
|
||||
assert runner._session_capacity == 1
|
||||
|
||||
runner._get_or_create_session("old")
|
||||
runner._get_or_create_session("new")
|
||||
|
||||
assert tuple(runner._sessions) == ("new",)
|
||||
assert pipeline.closes == ["old"]
|
||||
|
||||
|
||||
def test_dreamzero_like_requested_capacity_is_capped_by_budget():
|
||||
spec = dreamzero_like_spec(capacity=64)
|
||||
# Per all-layer self-KV page: 65,536 bytes. Two resident sessions need:
|
||||
# managed=(2 * (2 * 6 + 4) + 2)=34 pages, scratch=8 pages,
|
||||
# cross-attention=1 page/session, for 44 pages total.
|
||||
page_bytes = 65_536
|
||||
pipeline = CapablePipeline(spec)
|
||||
runner = make_runner(
|
||||
pipeline,
|
||||
available_bytes=44 * page_bytes,
|
||||
gpu_memory_fraction=1.0,
|
||||
)
|
||||
kv = runner.kv_cache
|
||||
assert kv is not None
|
||||
assert kv.requested_session_capacity == 64
|
||||
assert kv.session_capacity == 2
|
||||
assert runner._session_capacity == 2
|
||||
assert kv.cross_attention_reserved_bytes == 2 * page_bytes
|
||||
|
||||
|
||||
def test_forward_exception_releases_pending_allocation_and_model_state(monkeypatch):
|
||||
pipeline = CapablePipeline(lingbot_like_spec())
|
||||
runner = make_runner(pipeline)
|
||||
@@ -232,7 +367,10 @@ def test_forward_exception_releases_pending_allocation_and_model_state(monkeypat
|
||||
raise RuntimeError("layer exploded")
|
||||
|
||||
monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
|
||||
request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
|
||||
request = SimpleNamespace(
|
||||
request_id="broken-request",
|
||||
sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="layer exploded"):
|
||||
runner.execute_model(request)
|
||||
@@ -263,7 +401,10 @@ def test_synchronize_exception_uses_forward_cleanup_path(monkeypatch):
|
||||
monkeypatch.setattr(DiffusionModelRunner, "execute_model", return_after_allocation)
|
||||
monkeypatch.setattr(torch.accelerator, "synchronize", synchronize_boom)
|
||||
runner.device = torch.device("cuda")
|
||||
request = SimpleNamespace(sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}))
|
||||
request = SimpleNamespace(
|
||||
request_id="broken-request",
|
||||
sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="asynchronous kernel failed"):
|
||||
runner.execute_model(request)
|
||||
|
||||
@@ -0,0 +1,231 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator, Sequence
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.ar_diffusion.consumer import ARDiffusionOmniTickConsumer
|
||||
from vllm_omni.experimental.ar_diffusion.session import (
|
||||
ARDiffusionSession,
|
||||
ARDiffusionSessionEvent,
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.tick_protocol import (
|
||||
ARDiffusionChunkMetadata,
|
||||
ARDiffusionControlInput,
|
||||
ARDiffusionTickRequest,
|
||||
)
|
||||
from vllm_omni.inputs.data import (
|
||||
OmniDiffusionSamplingParams,
|
||||
OmniPromptType,
|
||||
OmniSamplingParams,
|
||||
)
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class FakeStagePool:
|
||||
def __init__(self, num_replicas: int = 1) -> None:
|
||||
self.num_replicas = num_replicas
|
||||
|
||||
|
||||
class FakeEngine:
|
||||
def __init__(self, num_replicas: int = 1) -> None:
|
||||
self.stage_pools = [FakeStagePool(num_replicas)]
|
||||
|
||||
|
||||
class FakeGenerateClient:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
include_metadata: bool = True,
|
||||
applied_event_ids: list[int] | None = None,
|
||||
output_error: str | None = None,
|
||||
num_replicas: int = 1,
|
||||
) -> None:
|
||||
self.engine = FakeEngine(num_replicas)
|
||||
self.include_metadata = include_metadata
|
||||
self.applied_event_ids = applied_event_ids
|
||||
self.output_error = output_error
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
prompt: OmniPromptType,
|
||||
*,
|
||||
request_id: str,
|
||||
sampling_params_list: Sequence[OmniSamplingParams],
|
||||
output_modalities: list[str] | None = None,
|
||||
) -> AsyncIterator[OmniRequestOutput]:
|
||||
self.calls.append(
|
||||
{
|
||||
"prompt": prompt,
|
||||
"request_id": request_id,
|
||||
"sampling_params_list": sampling_params_list,
|
||||
"output_modalities": output_modalities,
|
||||
}
|
||||
)
|
||||
params = sampling_params_list[0]
|
||||
assert isinstance(params, OmniDiffusionSamplingParams)
|
||||
tick = ARDiffusionTickRequest.from_extra_args(params.extra_args)
|
||||
assert tick is not None
|
||||
chunk_metadata = ARDiffusionChunkMetadata.from_tick(tick).to_dict()
|
||||
if self.applied_event_ids is not None:
|
||||
chunk_metadata["applied_event_ids"] = self.applied_event_ids
|
||||
metadata = {"metadata": {"ar_diffusion": chunk_metadata}} if self.include_metadata else {}
|
||||
yield OmniRequestOutput(
|
||||
request_id=request_id,
|
||||
stage_id=0,
|
||||
final_output_type="video",
|
||||
images=["chunk"], # type: ignore[list-item]
|
||||
_multimodal_output=metadata,
|
||||
finished=True,
|
||||
error=self.output_error,
|
||||
)
|
||||
|
||||
|
||||
class FakeLifecycle:
|
||||
async def reset_session(self, session_id: str) -> None:
|
||||
pass
|
||||
|
||||
async def close_session(self, session_id: str) -> None:
|
||||
pass
|
||||
|
||||
|
||||
def make_tick() -> ARDiffusionTickRequest:
|
||||
return ARDiffusionTickRequest(
|
||||
session_id="world-1",
|
||||
request_id="world-1-chunk-3",
|
||||
chunk_index=3,
|
||||
applied_event_ids=(7, 8),
|
||||
prompt="turn left",
|
||||
controls=(
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="example.camera.v1",
|
||||
data={"samples": [{"x": 1.0}]},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_consumer_rejects_replicated_ar_diffusion_stage() -> None:
|
||||
client = FakeGenerateClient(num_replicas=2)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=r"require exactly one replica.*num_replicas=2",
|
||||
):
|
||||
ARDiffusionOmniTickConsumer(
|
||||
client,
|
||||
prompt_provider=lambda tick: tick.prompt or "",
|
||||
sampling_params_list=[OmniDiffusionSamplingParams()],
|
||||
diffusion_stage_id=0,
|
||||
)
|
||||
|
||||
assert client.calls == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consumer_submits_typed_tick_and_returns_standard_output() -> None:
|
||||
client = FakeGenerateClient()
|
||||
template = OmniDiffusionSamplingParams(
|
||||
output_type="latent",
|
||||
extra_args={"deployment_option": "kept"},
|
||||
)
|
||||
consumer = ARDiffusionOmniTickConsumer(
|
||||
client,
|
||||
prompt_provider=lambda tick: {
|
||||
"prompt": tick.prompt,
|
||||
"multi_modal_data": {"image": "initial-image"},
|
||||
},
|
||||
sampling_params_list=[template],
|
||||
diffusion_stage_id=0,
|
||||
output_modalities=["video"],
|
||||
)
|
||||
tick = make_tick()
|
||||
|
||||
output = await consumer.execute_tick(tick)
|
||||
|
||||
assert output.images == ["chunk"]
|
||||
assert consumer.chunk_metadata(output) == ARDiffusionChunkMetadata.from_tick(tick)
|
||||
assert len(client.calls) == 1
|
||||
call = client.calls[0]
|
||||
assert call["request_id"] == tick.request_id
|
||||
assert call["output_modalities"] == ["video"]
|
||||
assert call["prompt"] == {
|
||||
"prompt": "turn left",
|
||||
"multi_modal_data": {"image": "initial-image"},
|
||||
}
|
||||
submitted_params = call["sampling_params_list"][0]
|
||||
assert submitted_params.extra_args["deployment_option"] == "kept"
|
||||
assert ARDiffusionTickRequest.from_extra_args(submitted_params.extra_args) == tick
|
||||
assert "ar_diffusion_tick" not in template.extra_args
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consumer_rejects_missing_chunk_metadata() -> None:
|
||||
consumer = ARDiffusionOmniTickConsumer(
|
||||
FakeGenerateClient(include_metadata=False),
|
||||
prompt_provider=lambda tick: tick.prompt or "",
|
||||
sampling_params_list=[OmniDiffusionSamplingParams()],
|
||||
diffusion_stage_id=0,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="multimodal_output.metadata"):
|
||||
await consumer.execute_tick(make_tick())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consumer_rejects_non_monotonic_applied_event_metadata() -> None:
|
||||
consumer = ARDiffusionOmniTickConsumer(
|
||||
FakeGenerateClient(applied_event_ids=[8, 7]),
|
||||
prompt_provider=lambda tick: tick.prompt or "",
|
||||
sampling_params_list=[OmniDiffusionSamplingParams()],
|
||||
diffusion_stage_id=0,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="unique and strictly increasing"):
|
||||
await consumer.execute_tick(make_tick())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_consumer_preserves_engine_failure_before_metadata_validation() -> None:
|
||||
consumer = ARDiffusionOmniTickConsumer(
|
||||
FakeGenerateClient(include_metadata=False, output_error="model failed"),
|
||||
prompt_provider=lambda tick: tick.prompt or "",
|
||||
sampling_params_list=[OmniDiffusionSamplingParams()],
|
||||
diffusion_stage_id=0,
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="model failed"):
|
||||
await consumer.execute_tick(make_tick())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_commits_only_after_standard_engine_output_metadata() -> None:
|
||||
client = FakeGenerateClient()
|
||||
consumer = ARDiffusionOmniTickConsumer(
|
||||
client,
|
||||
prompt_provider=lambda tick: tick.prompt or "",
|
||||
sampling_params_list=[OmniDiffusionSamplingParams()],
|
||||
diffusion_stage_id=0,
|
||||
)
|
||||
session = ARDiffusionSession(
|
||||
"world-1",
|
||||
tick_consumer=consumer,
|
||||
lifecycle=FakeLifecycle(),
|
||||
max_pending_events=4,
|
||||
request_id_factory=lambda session_id, chunk_index: f"{session_id}-chunk-{chunk_index}",
|
||||
)
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=7, prompt="turn left"))
|
||||
|
||||
output = await session.next_chunk()
|
||||
|
||||
assert output.request_id == "world-1-chunk-0"
|
||||
assert session.chunk_index == 1
|
||||
assert session.pending_event_count == 0
|
||||
assert client.calls[0]["request_id"] == "world-1-chunk-0"
|
||||
@@ -30,6 +30,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
BLOCK = 16
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False):
|
||||
return ChunkWindowSpec(
|
||||
@@ -95,7 +97,6 @@ def test_sliding_replace_snaps_to_chunk_boundary():
|
||||
def test_sink_chunks_protected():
|
||||
# sink = 1 chunk (16 tokens) is never skipped.
|
||||
skip = chunk_window_skipped_tokens(80, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=False)
|
||||
# raw overflow = 80 - 32 - 16 = 32 -> snapped 32
|
||||
assert skip == 32
|
||||
|
||||
|
||||
@@ -201,12 +202,149 @@ def test_cross_attn_pool_deducted_from_self_attn_budget():
|
||||
kv_branches=(ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 0)),
|
||||
session_capacity=1,
|
||||
cross_attention_lengths={"text": L},
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2 # K+V, pos+neg, layers
|
||||
expected = compute_num_blocks(avail - cross_bytes, 0.5, kv.spec.page_size_bytes * 2)
|
||||
page_bytes = kv.spec.page_size_bytes * 2
|
||||
scratch_bytes = kv.scratch_num_blocks * page_bytes
|
||||
expected = (int(avail * 0.5) - cross_bytes - scratch_bytes) // page_bytes
|
||||
assert kv.num_blocks == expected
|
||||
assert expected > 1 * (2 + 1) + 2 # above the local-slot floor, so the cross deduction is what's tested
|
||||
assert kv.cross_attention_reserved_bytes == cross_bytes
|
||||
assert (kv.num_blocks_total * page_bytes) + cross_bytes <= int(avail * 0.5)
|
||||
assert expected > 1 * (2 + 1) + 2 # above the local-slot minimum, so the cross deduction is what's tested
|
||||
|
||||
|
||||
def _make_tiny_capacity_kv(
|
||||
*,
|
||||
requested_capacity: int,
|
||||
available_bytes: int,
|
||||
gpu_memory_fraction: float = 1.0,
|
||||
model_owned_state_bytes_per_session: int = 0,
|
||||
) -> ARDiffusionKVCache:
|
||||
return ARDiffusionKVCache(
|
||||
ARDiffusionKVConfig(
|
||||
enable=True,
|
||||
chunk_size=1,
|
||||
window_chunks=3,
|
||||
sink_chunks=3,
|
||||
gpu_memory_fraction=gpu_memory_fraction,
|
||||
),
|
||||
num_layers=1,
|
||||
num_kv_heads=1,
|
||||
head_size=1,
|
||||
dtype=torch.float32,
|
||||
block_size=1,
|
||||
max_model_len=128,
|
||||
available_bytes=available_bytes,
|
||||
kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
|
||||
session_capacity=requested_capacity,
|
||||
cross_attention_lengths={"text": 2},
|
||||
frames_per_block=3,
|
||||
model_owned_state_bytes_per_session=model_owned_state_bytes_per_session,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
|
||||
def test_capacity_two_retains_both_windows_and_allocates_next_block():
|
||||
# LingBot-like single-branch geometry: page=8 bytes; capacity=2 requires
|
||||
# 17 managed + 3 scratch blocks, plus two 16-byte cross-attention
|
||||
# reservations: 192 bytes exactly.
|
||||
kv = _make_tiny_capacity_kv(requested_capacity=2, available_bytes=192)
|
||||
assert kv.requested_session_capacity == 2
|
||||
assert kv.session_capacity == 2
|
||||
assert kv.managed_num_blocks == 17
|
||||
|
||||
adapters = [kv.begin_request(session_id) for session_id in ("first", "second")]
|
||||
for adapter in adapters:
|
||||
for _ in range(6): # three sink + three window blocks
|
||||
kv.allocate_token_slots(adapter, 1)
|
||||
kv.commit_chunk(adapter)
|
||||
assert len(kv.window_block_ids(adapter)) == 6
|
||||
|
||||
# Only one request is in flight, so its three-frame block is counted once.
|
||||
kv.allocate_token_slots(adapters[0], 3)
|
||||
|
||||
|
||||
def test_requested_capacity_is_capped_and_cross_reservation_uses_effective_capacity():
|
||||
kv = _make_tiny_capacity_kv(requested_capacity=64, available_bytes=192)
|
||||
|
||||
assert kv.requested_session_capacity == 64
|
||||
assert kv.session_capacity == 2
|
||||
assert kv.cross_attention_bytes_per_session == 16
|
||||
assert kv.cross_attention_reserved_bytes == 32
|
||||
assert kv.num_blocks_total * 8 + kv.cross_attention_reserved_bytes == kv.memory_budget_bytes
|
||||
|
||||
|
||||
def test_capacity_rejects_budget_that_cannot_fit_one_session():
|
||||
# One session needs 11 managed + 3 scratch blocks and 16 cross bytes:
|
||||
# 14 * 8 + 16 = 128 bytes.
|
||||
with pytest.raises(ValueError, match="cannot fit one session.*available=127 bytes.*required=128 bytes"):
|
||||
_make_tiny_capacity_kv(requested_capacity=1, available_bytes=127)
|
||||
|
||||
|
||||
def test_fraction_is_soft_floor_when_one_session_fits_actual_memory():
|
||||
kv = _make_tiny_capacity_kv(
|
||||
requested_capacity=2,
|
||||
available_bytes=1280,
|
||||
gpu_memory_fraction=0.05,
|
||||
)
|
||||
|
||||
assert kv.configured_memory_budget_bytes == 64
|
||||
assert kv.memory_budget_bytes == 128
|
||||
assert kv.session_capacity == 1
|
||||
|
||||
|
||||
def test_model_owned_state_reduces_effective_session_capacity():
|
||||
kv = _make_tiny_capacity_kv(
|
||||
requested_capacity=2,
|
||||
available_bytes=200,
|
||||
model_owned_state_bytes_per_session=16,
|
||||
)
|
||||
|
||||
assert kv.session_capacity == 1
|
||||
assert kv.model_owned_state_reserved_bytes == 16
|
||||
assert (
|
||||
kv.num_blocks_total * 8 + kv.cross_attention_reserved_bytes + kv.model_owned_state_reserved_bytes
|
||||
<= kv.memory_budget_bytes
|
||||
)
|
||||
|
||||
|
||||
def _make_shipped_lingbot_geometry(*, gpu_memory_fraction: float) -> ARDiffusionKVCache:
|
||||
return ARDiffusionKVCache(
|
||||
ARDiffusionKVConfig(
|
||||
enable=True,
|
||||
chunk_size=1560,
|
||||
window_chunks=9,
|
||||
sink_chunks=9,
|
||||
gpu_memory_fraction=gpu_memory_fraction,
|
||||
),
|
||||
num_layers=40,
|
||||
num_kv_heads=40,
|
||||
head_size=128,
|
||||
dtype=torch.bfloat16,
|
||||
block_size=1560,
|
||||
max_model_len=1 << 20,
|
||||
available_bytes=100 * (1 << 30),
|
||||
kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
|
||||
session_capacity=2,
|
||||
cross_attention_lengths={"text": 512},
|
||||
frames_per_block=3,
|
||||
model_owned_state_bytes_per_session=7_488_000,
|
||||
)
|
||||
|
||||
|
||||
def test_shipped_lingbot_geometry_default_fraction_admits_one_session():
|
||||
kv = _make_shipped_lingbot_geometry(gpu_memory_fraction=0.1)
|
||||
|
||||
assert kv.configured_memory_budget_bytes == 10 * (1 << 30)
|
||||
assert kv.memory_budget_bytes == 33_653_670_400
|
||||
assert kv.session_capacity == 1
|
||||
|
||||
|
||||
def test_shipped_lingbot_geometry_tuned_fraction_admits_two_sessions():
|
||||
kv = _make_shipped_lingbot_geometry(gpu_memory_fraction=0.6)
|
||||
|
||||
assert kv.memory_budget_bytes == 60 * (1 << 30)
|
||||
assert kv.session_capacity == 2
|
||||
|
||||
|
||||
def _make_kv(
|
||||
@@ -221,15 +359,24 @@ def _make_kv(
|
||||
if local_branches == 1
|
||||
else (ARDiffusionKVBranchSpec("positive", 0), ARDiffusionKVBranchSpec("negative", 1))
|
||||
)
|
||||
page_bytes = 2 * BLOCK * 4 * 64 * torch.float32.itemsize
|
||||
declared_scratch_blocks = (max_scratch_tokens_per_branch + BLOCK - 1) // BLOCK
|
||||
scratch_blocks = local_branches * (num_frame_per_block + declared_scratch_blocks)
|
||||
managed_blocks = local_branches * (window_chunks + num_frame_per_block) + 2
|
||||
return ARDiffusionKVCache(
|
||||
ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks),
|
||||
ARDiffusionKVConfig(
|
||||
enable=True,
|
||||
chunk_size=BLOCK,
|
||||
window_chunks=window_chunks,
|
||||
gpu_memory_fraction=1.0,
|
||||
),
|
||||
num_layers=1,
|
||||
num_kv_heads=4,
|
||||
head_size=64,
|
||||
dtype=torch.float32,
|
||||
block_size=BLOCK,
|
||||
max_model_len=4096,
|
||||
available_bytes=1 << 16, # tiny -> the floor binds
|
||||
available_bytes=(managed_blocks + scratch_blocks) * page_bytes,
|
||||
device=torch.device("cpu"),
|
||||
kv_branches=kv_branches,
|
||||
session_capacity=1,
|
||||
|
||||
@@ -54,13 +54,15 @@ def test_full_request_lifecycle_and_eviction():
|
||||
used = {int(s) // kv.block_size for s in slots}
|
||||
assert kv.null_block_id not in used
|
||||
assert used <= set(block_table)
|
||||
# Resident window stays bounded by window + the in-flight chunk.
|
||||
# Allocation may transiently add one in-flight chunk.
|
||||
assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks + 1
|
||||
free_after.append(kv.manager.block_pool.get_num_free_blocks())
|
||||
kv.commit_chunk(adapter)
|
||||
# A successful commit immediately prunes physical ownership.
|
||||
assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks
|
||||
free_after.append(kv.manager.block_pool.get_num_free_blocks())
|
||||
|
||||
# Pool memory plateaus once the window is full (eviction recycles blocks).
|
||||
assert free_after[-1] == free_after[kv.spec.window_chunks]
|
||||
assert free_after[-1] == free_after[kv.spec.window_chunks - 1]
|
||||
|
||||
kv.end_request(adapter)
|
||||
assert kv.manager.block_pool.get_num_free_blocks() == free_total
|
||||
|
||||
@@ -60,9 +60,18 @@ def _dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
|
||||
return torch.einsum("bhqk,bkhd->bqhd", probs, value)
|
||||
|
||||
|
||||
def _cuda_flash_attn_usable() -> bool:
|
||||
def _gpu_flash_attn_usable() -> bool:
|
||||
if not torch.cuda.is_available():
|
||||
return False
|
||||
if torch.version.hip is not None:
|
||||
for module in ("aiter", "flash_attn"):
|
||||
try:
|
||||
imported = __import__(module, fromlist=["flash_attn_varlen_func"])
|
||||
if getattr(imported, "flash_attn_varlen_func", None) is not None:
|
||||
return True
|
||||
except ImportError:
|
||||
pass
|
||||
return False
|
||||
try:
|
||||
spec = find_spec("vllm.vllm_flash_attn")
|
||||
if spec is None or spec.origin is None:
|
||||
@@ -227,12 +236,11 @@ def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len,
|
||||
assert st.adapter(POS).completed_chunks == before + (1 if commit_current else 0)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not _cuda_flash_attn_usable(), reason="usable CUDA FlashAttention is required")
|
||||
@pytest.mark.skipif(not _gpu_flash_attn_usable(), reason="usable GPU FlashAttention is required")
|
||||
@pytest.mark.parametrize("history_chunks", [1, 3])
|
||||
@pytest.mark.parametrize("action_len", [0, 3])
|
||||
@pytest.mark.parametrize("commit_current", [False, True])
|
||||
def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len, commit_current):
|
||||
pytest.importorskip("vllm.vllm_flash_attn")
|
||||
torch.manual_seed(0)
|
||||
device = torch.device("cuda")
|
||||
dtype = torch.float16
|
||||
@@ -324,7 +332,7 @@ def test_block_table_padded_to_fixed_width():
|
||||
|
||||
def test_prepare_is_idempotent_and_layers_share_metadata():
|
||||
device = torch.device("cpu")
|
||||
_, st = make_state(num_layers=2)
|
||||
kv, st = make_state(num_layers=2)
|
||||
contexts = st.get_kv_caches(POS, seq_len=BLOCK, commit_current=False)
|
||||
fctx = contexts[0].forward_ctx
|
||||
fctx.prepare(device=device, action_len=0, query_len=BLOCK)
|
||||
@@ -341,6 +349,10 @@ def test_prepare_is_idempotent_and_layers_share_metadata():
|
||||
assert i0.block_table is i1.block_table
|
||||
assert i0.seq_lens is i1.seq_lens
|
||||
assert i0.video_slots is i1.video_slots
|
||||
assert i0.key_pool is kv._k_pools[0]
|
||||
assert i0.value_pool is kv._v_pools[0]
|
||||
assert i1.key_pool is kv._k_pools[1]
|
||||
assert i1.value_pool is kv._v_pools[1]
|
||||
|
||||
|
||||
def test_layer_inputs_before_prepare_raises():
|
||||
|
||||
@@ -174,6 +174,19 @@ def test_window_reset_can_keep_named_cross_attention_cache():
|
||||
assert not st.is_cross_attention_populated(POS, "text")
|
||||
|
||||
|
||||
def test_clear_cross_attention_preserves_self_attention_history():
|
||||
kv, st = make_state(num_layers=1, cross_attn_length=8)
|
||||
_prepare_and_commit(st, POS, 2)
|
||||
k = torch.randn(1, 8, N_HEADS, HEAD_DIM)
|
||||
st.populate_cross_attention(POS, "text", [(k, k)])
|
||||
resident_before = kv.window_block_ids(st.adapter(POS))
|
||||
|
||||
st.clear_cross_attention()
|
||||
|
||||
assert kv.window_block_ids(st.adapter(POS)) == resident_before
|
||||
assert not st.is_cross_attention_populated(POS, "text")
|
||||
|
||||
|
||||
def test_cross_attention_uses_logical_branch_when_local_slot_is_shared():
|
||||
_, st = make_state(num_layers=2, cross_attn_length=8, shared_local_index=True)
|
||||
assert st.kv_cache.num_local_kv_branches == 1
|
||||
|
||||
@@ -0,0 +1,505 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.ar_diffusion.session import (
|
||||
ARDiffusionChunkMetadataError,
|
||||
ARDiffusionEventAcceptanceStatus,
|
||||
ARDiffusionEventQueueFullError,
|
||||
ARDiffusionPreparedControls,
|
||||
ARDiffusionSession,
|
||||
ARDiffusionSessionError,
|
||||
ARDiffusionSessionEvent,
|
||||
ARDiffusionSessionManager,
|
||||
ARDiffusionSessionStateError,
|
||||
ARDiffusionSessionStatus,
|
||||
ARDiffusionStaleEventError,
|
||||
ARDiffusionWorkerLifecycle,
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.tick_protocol import (
|
||||
ARDiffusionChunkMetadata,
|
||||
ARDiffusionControlInput,
|
||||
ARDiffusionTickRequest,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class FakeTickConsumer:
|
||||
def __init__(self) -> None:
|
||||
self.ticks: list[ARDiffusionTickRequest] = []
|
||||
self.failures_remaining = 0
|
||||
self.started: asyncio.Event | None = None
|
||||
self.release: asyncio.Event | None = None
|
||||
self.metadata_override: ARDiffusionChunkMetadata | None = None
|
||||
|
||||
async def execute_tick(self, tick: ARDiffusionTickRequest) -> ARDiffusionChunkMetadata:
|
||||
self.ticks.append(tick)
|
||||
if self.started is not None:
|
||||
self.started.set()
|
||||
if self.release is not None:
|
||||
await self.release.wait()
|
||||
if self.failures_remaining:
|
||||
self.failures_remaining -= 1
|
||||
raise RuntimeError("fake tick failed")
|
||||
return self.metadata_override or ARDiffusionChunkMetadata.from_tick(tick)
|
||||
|
||||
def chunk_metadata(self, output: ARDiffusionChunkMetadata) -> ARDiffusionChunkMetadata:
|
||||
return output
|
||||
|
||||
|
||||
class FakeLifecycle:
|
||||
def __init__(self) -> None:
|
||||
self.resets: list[str] = []
|
||||
self.closes: list[str] = []
|
||||
self.close_failures_remaining = 0
|
||||
|
||||
async def reset_session(self, session_id: str) -> None:
|
||||
self.resets.append(session_id)
|
||||
|
||||
async def close_session(self, session_id: str) -> None:
|
||||
self.closes.append(session_id)
|
||||
if self.close_failures_remaining:
|
||||
self.close_failures_remaining -= 1
|
||||
raise RuntimeError("fake close failed")
|
||||
|
||||
|
||||
class FailingCommitReducer:
|
||||
def __init__(self) -> None:
|
||||
self.commit_calls = 0
|
||||
self.reset_calls = 0
|
||||
|
||||
def prepare(
|
||||
self,
|
||||
*,
|
||||
current_controls: dict[str, ARDiffusionControlInput],
|
||||
events: tuple[ARDiffusionSessionEvent, ...],
|
||||
chunk_index: int,
|
||||
) -> ARDiffusionPreparedControls:
|
||||
del current_controls, events, chunk_index
|
||||
return ARDiffusionPreparedControls(controls=())
|
||||
|
||||
def commit(self, prepared: ARDiffusionPreparedControls) -> None:
|
||||
del prepared
|
||||
self.commit_calls += 1
|
||||
raise RuntimeError("fake reducer commit failed")
|
||||
|
||||
def reset(self) -> None:
|
||||
self.reset_calls += 1
|
||||
|
||||
|
||||
def make_session(
|
||||
*,
|
||||
max_pending_events: int = 8,
|
||||
consumer: FakeTickConsumer | None = None,
|
||||
lifecycle: FakeLifecycle | None = None,
|
||||
) -> tuple[ARDiffusionSession, FakeTickConsumer, FakeLifecycle]:
|
||||
resolved_consumer = consumer or FakeTickConsumer()
|
||||
resolved_lifecycle = lifecycle or FakeLifecycle()
|
||||
session = ARDiffusionSession(
|
||||
"world-1",
|
||||
tick_consumer=resolved_consumer,
|
||||
lifecycle=resolved_lifecycle,
|
||||
max_pending_events=max_pending_events,
|
||||
request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
|
||||
)
|
||||
return session, resolved_consumer, resolved_lifecycle
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_events_are_accepted_then_sorted_and_applied_at_chunk_boundary() -> None:
|
||||
session, consumer, _ = make_session()
|
||||
event_two = ARDiffusionSessionEvent(event_id=2, prompt="turn")
|
||||
event_one = ARDiffusionSessionEvent(
|
||||
event_id=1,
|
||||
controls=(
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="example.camera.v1",
|
||||
data={"samples": [{"x": 1.0}]},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
accepted_two = await session.accept_event(event_two)
|
||||
accepted_one = await session.accept_event(event_one)
|
||||
duplicate = await session.accept_event(event_two)
|
||||
|
||||
assert accepted_two.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
|
||||
assert accepted_one.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
|
||||
assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
|
||||
assert consumer.ticks == []
|
||||
|
||||
metadata = await session.next_chunk()
|
||||
|
||||
assert metadata.applied_event_ids == (1, 2)
|
||||
assert consumer.ticks[0] == ARDiffusionTickRequest(
|
||||
session_id="world-1",
|
||||
request_id="world-1-request-0",
|
||||
chunk_index=0,
|
||||
applied_event_ids=(1, 2),
|
||||
prompt="turn",
|
||||
controls=(
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="example.camera.v1",
|
||||
data={"samples": [{"x": 1.0}]},
|
||||
),
|
||||
),
|
||||
)
|
||||
assert session.pending_event_count == 0
|
||||
|
||||
applied_duplicate = await session.accept_event(event_two)
|
||||
assert applied_duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
|
||||
|
||||
next_metadata = await session.next_chunk()
|
||||
assert next_metadata.chunk_index == 1
|
||||
assert next_metadata.applied_event_ids == ()
|
||||
assert consumer.ticks[1].prompt == "turn"
|
||||
assert consumer.ticks[1].controls == consumer.ticks[0].controls
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_composite_event_is_isolated_and_queue_admission_is_atomic() -> None:
|
||||
session, consumer, _ = make_session(max_pending_events=1)
|
||||
payload = {"samples": [{"x": 1.0}]}
|
||||
composite = ARDiffusionSessionEvent(
|
||||
event_id=5,
|
||||
prompt="new prompt",
|
||||
controls=(
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="example.camera.v1",
|
||||
data=payload,
|
||||
),
|
||||
ARDiffusionControlInput(
|
||||
track="action",
|
||||
schema="example.action.v1",
|
||||
data={"keys": ["left"]},
|
||||
),
|
||||
),
|
||||
)
|
||||
await session.accept_event(composite)
|
||||
payload["samples"][0]["x"] = 99.0
|
||||
|
||||
duplicate = await session.accept_event(composite)
|
||||
assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
|
||||
with pytest.raises(ARDiffusionEventQueueFullError, match="queue is full"):
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=6, prompt="rejected"))
|
||||
|
||||
await session.next_chunk()
|
||||
assert consumer.ticks[0].applied_event_ids == (5,)
|
||||
assert consumer.ticks[0].prompt == "new prompt"
|
||||
assert consumer.ticks[0].controls[0].data["keys"] == ("left",)
|
||||
assert consumer.ticks[0].controls[1].data["samples"] == ({"x": 1.0},)
|
||||
assert consumer.ticks[0].controls[0].to_dict()["data"] == {"keys": ["left"]}
|
||||
assert consumer.ticks[0].controls[1].to_dict()["data"] == {"samples": [{"x": 1.0}]}
|
||||
|
||||
|
||||
def test_event_rejects_duplicate_tracks_and_non_transport_payloads() -> None:
|
||||
with pytest.raises(ValueError, match="at most once"):
|
||||
ARDiffusionSessionEvent(
|
||||
event_id=1,
|
||||
controls=(
|
||||
ARDiffusionControlInput(track="camera", schema="one", data={}),
|
||||
ARDiffusionControlInput(track="camera", schema="two", data={}),
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="transport-safe"):
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="example.camera.v1",
|
||||
data={"object": object()},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_tick_closes_worker_and_requires_reset() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
consumer.failures_remaining = 1
|
||||
lifecycle = FakeLifecycle()
|
||||
session, _, _ = make_session(consumer=consumer, lifecycle=lifecycle)
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=3, prompt="retry me"))
|
||||
|
||||
with pytest.raises(RuntimeError, match="fake tick failed"):
|
||||
await session.next_chunk()
|
||||
|
||||
assert session.status is ARDiffusionSessionStatus.FAILED
|
||||
assert session.chunk_index == 0
|
||||
assert session.pending_event_count == 1
|
||||
assert lifecycle.closes == ["world-1"]
|
||||
with pytest.raises(ARDiffusionSessionStateError, match="failed"):
|
||||
await session.next_chunk()
|
||||
|
||||
await session.reset()
|
||||
assert session.status is ARDiffusionSessionStatus.ACTIVE
|
||||
assert session.pending_event_count == 0
|
||||
assert lifecycle.resets == ["world-1"]
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="after reset"))
|
||||
metadata = await session.next_chunk()
|
||||
assert metadata.chunk_index == 0
|
||||
assert metadata.applied_event_ids == (4,)
|
||||
assert [tick.applied_event_ids for tick in consumer.ticks] == [(3,), (4,)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inflight_snapshot_is_immutable_and_new_event_waits_for_next_chunk() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
consumer.started = asyncio.Event()
|
||||
consumer.release = asyncio.Event()
|
||||
session, _, _ = make_session(consumer=consumer)
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=10, prompt="first"))
|
||||
|
||||
first_chunk = asyncio.create_task(session.next_chunk())
|
||||
await consumer.started.wait()
|
||||
accepted = await session.accept_event(ARDiffusionSessionEvent(event_id=11, prompt="second"))
|
||||
assert accepted.status is ARDiffusionEventAcceptanceStatus.ACCEPTED
|
||||
assert consumer.ticks[0].prompt == "first"
|
||||
assert consumer.ticks[0].applied_event_ids == (10,)
|
||||
|
||||
with pytest.raises(ARDiffusionStaleEventError, match="event floor is 10"):
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=9, prompt="stale"))
|
||||
|
||||
consumer.release.set()
|
||||
first_metadata = await first_chunk
|
||||
second_metadata = await session.next_chunk()
|
||||
assert first_metadata.applied_event_ids == (10,)
|
||||
assert second_metadata.applied_event_ids == (11,)
|
||||
assert consumer.ticks[1].prompt == "second"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metadata_mismatch_does_not_commit_snapshot() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
consumer.metadata_override = ARDiffusionChunkMetadata(
|
||||
session_id="another-world",
|
||||
request_id="wrong",
|
||||
chunk_index=0,
|
||||
)
|
||||
lifecycle = FakeLifecycle()
|
||||
session, _, _ = make_session(consumer=consumer, lifecycle=lifecycle)
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="pending"))
|
||||
|
||||
with pytest.raises(ARDiffusionChunkMetadataError, match="does not match"):
|
||||
await session.next_chunk()
|
||||
|
||||
assert session.chunk_index == 0
|
||||
assert session.pending_event_count == 1
|
||||
assert session.status is ARDiffusionSessionStatus.FAILED
|
||||
assert lifecycle.closes == ["world-1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reducer_commit_failure_is_fail_closed_without_partial_session_commit() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
lifecycle = FakeLifecycle()
|
||||
reducer = FailingCommitReducer()
|
||||
session = ARDiffusionSession(
|
||||
"world-1",
|
||||
tick_consumer=consumer,
|
||||
lifecycle=lifecycle,
|
||||
max_pending_events=4,
|
||||
request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
|
||||
control_reducer=reducer,
|
||||
)
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="pending"))
|
||||
|
||||
with pytest.raises(RuntimeError, match="fake reducer commit failed"):
|
||||
await session.next_chunk()
|
||||
|
||||
assert reducer.commit_calls == 1
|
||||
assert session.status is ARDiffusionSessionStatus.FAILED
|
||||
assert session.chunk_index == 0
|
||||
assert session.pending_event_count == 1
|
||||
assert lifecycle.closes == ["world-1"]
|
||||
with pytest.raises(ARDiffusionSessionStateError, match="failed"):
|
||||
await session.next_chunk()
|
||||
|
||||
await session.reset()
|
||||
assert reducer.reset_calls == 1
|
||||
assert session.status is ARDiffusionSessionStatus.ACTIVE
|
||||
assert session.chunk_index == 0
|
||||
assert session.pending_event_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_close_and_disconnect_reach_lifecycle_boundary() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
lifecycle = FakeLifecycle()
|
||||
manager = ARDiffusionSessionManager(
|
||||
tick_consumer=consumer,
|
||||
lifecycle=lifecycle,
|
||||
max_pending_events=4,
|
||||
request_id_factory=lambda session_id, chunk_index: f"{session_id}-{chunk_index}",
|
||||
)
|
||||
session = await manager.create_session("world-reset")
|
||||
await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="before reset"))
|
||||
await session.next_chunk()
|
||||
await manager.reset_session("world-reset")
|
||||
|
||||
assert lifecycle.resets == ["world-reset"]
|
||||
assert session.status is ARDiffusionSessionStatus.ACTIVE
|
||||
assert session.chunk_index == 0
|
||||
duplicate = await session.accept_event(ARDiffusionSessionEvent(event_id=4, prompt="duplicate"))
|
||||
assert duplicate.status is ARDiffusionEventAcceptanceStatus.DUPLICATE
|
||||
|
||||
reset_tick = await session.next_chunk()
|
||||
assert reset_tick.chunk_index == 0
|
||||
assert consumer.ticks[-1].prompt is None
|
||||
assert consumer.ticks[-1].controls == ()
|
||||
|
||||
await manager.close_session("world-reset")
|
||||
assert lifecycle.closes == ["world-reset"]
|
||||
assert session.status is ARDiffusionSessionStatus.CLOSED
|
||||
with pytest.raises(ARDiffusionSessionStateError, match="does not exist"):
|
||||
await manager.get_session("world-reset")
|
||||
|
||||
disconnected = await manager.create_session("world-disconnect")
|
||||
await manager.disconnect("world-disconnect")
|
||||
assert disconnected.status is ARDiffusionSessionStatus.CLOSED
|
||||
assert lifecycle.closes[-1] == "world-disconnect"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_manager_keeps_interleaved_sessions_independent() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
manager = ARDiffusionSessionManager(
|
||||
tick_consumer=consumer,
|
||||
lifecycle=FakeLifecycle(),
|
||||
max_pending_events=4,
|
||||
request_id_factory=lambda session_id, chunk_index: f"{session_id}-request-{chunk_index}",
|
||||
)
|
||||
session_a = await manager.create_session("world-a")
|
||||
session_b = await manager.create_session("world-b")
|
||||
|
||||
await session_a.accept_event(ARDiffusionSessionEvent(event_id=10, prompt="a0"))
|
||||
await session_a.next_chunk()
|
||||
await session_b.accept_event(ARDiffusionSessionEvent(event_id=1, prompt="b0"))
|
||||
await session_b.next_chunk()
|
||||
await session_a.accept_event(ARDiffusionSessionEvent(event_id=11, prompt="a1"))
|
||||
await session_a.next_chunk()
|
||||
|
||||
assert [
|
||||
(
|
||||
tick.session_id,
|
||||
tick.request_id,
|
||||
tick.chunk_index,
|
||||
tick.applied_event_ids,
|
||||
)
|
||||
for tick in consumer.ticks
|
||||
] == [
|
||||
("world-a", "world-a-request-0", 0, (10,)),
|
||||
("world-b", "world-b-request-0", 0, (1,)),
|
||||
("world-a", "world-a-request-1", 1, (11,)),
|
||||
]
|
||||
assert session_a.chunk_index == 2
|
||||
assert session_b.chunk_index == 1
|
||||
with pytest.raises(ARDiffusionStaleEventError, match="event floor is 11"):
|
||||
await session_a.accept_event(ARDiffusionSessionEvent(event_id=2, prompt="stale only for a"))
|
||||
assert (
|
||||
await session_b.accept_event(ARDiffusionSessionEvent(event_id=2, prompt="b1"))
|
||||
).status is ARDiffusionEventAcceptanceStatus.ACCEPTED
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_close_retains_tombstone_until_cleanup_retry_succeeds() -> None:
|
||||
consumer = FakeTickConsumer()
|
||||
lifecycle = FakeLifecycle()
|
||||
lifecycle.close_failures_remaining = 1
|
||||
manager = ARDiffusionSessionManager(
|
||||
tick_consumer=consumer,
|
||||
lifecycle=lifecycle,
|
||||
max_pending_events=4,
|
||||
)
|
||||
session = await manager.create_session("world-close-retry")
|
||||
|
||||
with pytest.raises(RuntimeError, match="fake close failed"):
|
||||
await manager.close_session("world-close-retry")
|
||||
|
||||
assert session.status is ARDiffusionSessionStatus.CLEANUP_FAILED
|
||||
assert await manager.get_session("world-close-retry") is session
|
||||
with pytest.raises(ARDiffusionSessionStateError, match="already exists"):
|
||||
await manager.create_session("world-close-retry")
|
||||
with pytest.raises(ARDiffusionSessionStateError, match="cleanup_failed"):
|
||||
await manager.reset_session("world-close-retry")
|
||||
|
||||
await manager.close_session("world-close-retry")
|
||||
|
||||
assert lifecycle.closes == ["world-close-retry", "world-close-retry"]
|
||||
assert session.status is ARDiffusionSessionStatus.CLOSED
|
||||
with pytest.raises(ARDiffusionSessionStateError, match="does not exist"):
|
||||
await manager.get_session("world-close-retry")
|
||||
recreated = await manager.create_session("world-close-retry")
|
||||
assert recreated is not session
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_disconnect_retains_session_for_explicit_cleanup_retry() -> None:
|
||||
lifecycle = FakeLifecycle()
|
||||
lifecycle.close_failures_remaining = 1
|
||||
manager = ARDiffusionSessionManager(
|
||||
tick_consumer=FakeTickConsumer(),
|
||||
lifecycle=lifecycle,
|
||||
max_pending_events=4,
|
||||
)
|
||||
session = await manager.create_session("world-disconnect-retry")
|
||||
|
||||
with pytest.raises(RuntimeError, match="fake close failed"):
|
||||
await manager.disconnect("world-disconnect-retry")
|
||||
|
||||
assert session.status is ARDiffusionSessionStatus.CLEANUP_FAILED
|
||||
assert await manager.get_session("world-disconnect-retry") is session
|
||||
|
||||
await manager.close_session("world-disconnect-retry")
|
||||
assert session.status is ARDiffusionSessionStatus.CLOSED
|
||||
|
||||
|
||||
class FakeRPCClient:
|
||||
def __init__(self, result: list[Any]) -> None:
|
||||
self.result = result
|
||||
self.calls: list[dict[str, Any]] = []
|
||||
|
||||
async def collective_rpc(self, **kwargs: Any) -> list[Any]:
|
||||
self.calls.append(kwargs)
|
||||
return self.result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_worker_lifecycle_routes_to_selected_stage_rpc() -> None:
|
||||
client = FakeRPCClient([[True]])
|
||||
lifecycle = ARDiffusionWorkerLifecycle(client, stage_ids=[2], timeout=3.0)
|
||||
|
||||
await lifecycle.reset_session("world-1")
|
||||
await lifecycle.close_session("world-1")
|
||||
|
||||
assert client.calls == [
|
||||
{
|
||||
"method": "reset_ar_diffusion_session",
|
||||
"timeout": 3.0,
|
||||
"args": ("world-1",),
|
||||
"stage_ids": [2],
|
||||
},
|
||||
{
|
||||
"method": "close_ar_diffusion_session",
|
||||
"timeout": 3.0,
|
||||
"args": ("world-1",),
|
||||
"stage_ids": [2],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_worker_lifecycle_rejects_unsupported_stage_result() -> None:
|
||||
lifecycle = ARDiffusionWorkerLifecycle(
|
||||
FakeRPCClient([{"supported": False, "error": "not an AR stage"}]),
|
||||
stage_ids=[1],
|
||||
)
|
||||
|
||||
with pytest.raises(ARDiffusionSessionError, match="not an AR stage"):
|
||||
await lifecycle.close_session("world-1")
|
||||
@@ -0,0 +1,108 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.experimental.ar_diffusion.tick_protocol import (
|
||||
ARDiffusionChunkMetadata,
|
||||
ARDiffusionControlInput,
|
||||
ARDiffusionTickRequest,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_tick_round_trip_preserves_identity_and_control_snapshot() -> None:
|
||||
tick = ARDiffusionTickRequest(
|
||||
session_id="world-1",
|
||||
request_id="chunk-request-7",
|
||||
chunk_index=7,
|
||||
applied_event_ids=(10, 11),
|
||||
prompt="turn left",
|
||||
controls=(
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="lingbot.camera_trajectory.v1",
|
||||
data={"poses": [[1.0, 0.0]], "intrinsics": [[2.0]]},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
restored = ARDiffusionTickRequest.from_extra_args(tick.to_extra_args())
|
||||
|
||||
assert restored == tick
|
||||
assert ARDiffusionChunkMetadata.from_tick(tick).to_dict() == {
|
||||
"session_id": "world-1",
|
||||
"request_id": "chunk-request-7",
|
||||
"chunk_index": 7,
|
||||
"applied_event_ids": [10, 11],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("field", "value", "message"),
|
||||
[
|
||||
("session_id", "", "session_id"),
|
||||
("request_id", "", "request_id"),
|
||||
("chunk_index", -1, "chunk_index"),
|
||||
("applied_event_ids", (2, 1), "strictly increasing"),
|
||||
("applied_event_ids", (1, 1), "strictly increasing"),
|
||||
],
|
||||
)
|
||||
def test_tick_rejects_invalid_identity(
|
||||
field: str,
|
||||
value: object,
|
||||
message: str,
|
||||
) -> None:
|
||||
values = {
|
||||
"session_id": "world-1",
|
||||
"request_id": "request-1",
|
||||
"chunk_index": 0,
|
||||
"applied_event_ids": (),
|
||||
}
|
||||
values[field] = value
|
||||
|
||||
with pytest.raises(ValueError, match=message):
|
||||
ARDiffusionTickRequest(**values)
|
||||
|
||||
|
||||
def test_tick_snapshot_is_detached_and_deeply_immutable() -> None:
|
||||
payload = {"samples": [{"keys": ["w"]}]}
|
||||
event_ids = [1, 2]
|
||||
controls = [
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="example.camera.v1",
|
||||
data=payload,
|
||||
)
|
||||
]
|
||||
tick = ARDiffusionTickRequest(
|
||||
session_id="world-1",
|
||||
request_id="request-a",
|
||||
chunk_index=0,
|
||||
applied_event_ids=event_ids, # type: ignore[arg-type]
|
||||
controls=controls, # type: ignore[arg-type]
|
||||
)
|
||||
serialized = tick.to_extra_args()
|
||||
|
||||
payload["samples"][0]["keys"].append("j")
|
||||
event_ids.append(3)
|
||||
controls.clear()
|
||||
serialized["ar_diffusion_tick"]["controls"][0]["data"]["samples"][0]["keys"].append("k")
|
||||
|
||||
assert tick.applied_event_ids == (1, 2)
|
||||
assert isinstance(tick.controls, tuple)
|
||||
assert tick.to_extra_args()["ar_diffusion_tick"]["controls"] == [
|
||||
{
|
||||
"track": "camera",
|
||||
"schema": "example.camera.v1",
|
||||
"data": {"samples": [{"keys": ["w"]}]},
|
||||
}
|
||||
]
|
||||
with pytest.raises(TypeError):
|
||||
tick.controls[0].data["samples"] = () # type: ignore[index]
|
||||
with pytest.raises(AttributeError):
|
||||
tick.controls[0].data["samples"][0]["keys"].append("j")
|
||||
|
||||
|
||||
def test_legacy_request_has_no_typed_tick() -> None:
|
||||
assert ARDiffusionTickRequest.from_extra_args({"session_id": "legacy"}) is None
|
||||
@@ -0,0 +1,53 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class FakeARDiffusionRunner:
|
||||
def __init__(self) -> None:
|
||||
self.resets: list[str] = []
|
||||
self.closes: list[str] = []
|
||||
|
||||
def reset_session(self, session_id: str) -> None:
|
||||
self.resets.append(session_id)
|
||||
|
||||
def close_session(self, session_id: str) -> None:
|
||||
self.closes.append(session_id)
|
||||
|
||||
|
||||
class FakeRegularRunner:
|
||||
pass
|
||||
|
||||
|
||||
def make_worker(runner: object) -> DiffusionWorker:
|
||||
worker = object.__new__(DiffusionWorker)
|
||||
worker.model_runner = runner
|
||||
return worker
|
||||
|
||||
|
||||
def test_worker_delegates_ar_session_lifecycle_to_runner() -> None:
|
||||
runner = FakeARDiffusionRunner()
|
||||
worker = make_worker(runner)
|
||||
|
||||
assert worker.reset_ar_diffusion_session("world-1") is True
|
||||
assert worker.close_ar_diffusion_session("world-1") is True
|
||||
assert runner.resets == ["world-1"]
|
||||
assert runner.closes == ["world-1"]
|
||||
|
||||
|
||||
def test_worker_reports_unsupported_runner_without_model_assumptions() -> None:
|
||||
worker = make_worker(FakeRegularRunner())
|
||||
|
||||
assert worker.reset_ar_diffusion_session("world-1") is False
|
||||
assert worker.close_ar_diffusion_session("world-1") is False
|
||||
|
||||
|
||||
def test_worker_rejects_empty_session_id() -> None:
|
||||
worker = make_worker(FakeARDiffusionRunner())
|
||||
|
||||
with pytest.raises(ValueError, match="session_id"):
|
||||
worker.close_ar_diffusion_session("")
|
||||
@@ -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 per-role attention backend configuration (RFC: per-role-attention-backend).
|
||||
|
||||
@@ -90,10 +90,23 @@ class TestAttentionSpec:
|
||||
with pytest.raises(ValueError, match=match):
|
||||
AttentionSpec(**spec)
|
||||
|
||||
def test_fastvideo_vsa_topk_serialized(self):
|
||||
spec = AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=96)
|
||||
assert spec.backend_kwargs() == {"topk": 96}
|
||||
|
||||
def test_fastvideo_vsa_topk_rejected_for_other_backend(self):
|
||||
with pytest.raises(ValueError, match="only supported by the FASTVIDEO_VSA"):
|
||||
AttentionSpec(backend="TORCH_SDPA", fastvideo_vsa_topk=96)
|
||||
|
||||
def test_block_sparse_defaults_applied_when_backend_selected(self):
|
||||
spec = AttentionSpec(backend="RAINFUSION_ATTN")
|
||||
assert spec.block_sparse.sparsity == 0.8
|
||||
assert spec.backend_kwargs() == {"sparsity": 0.8, "start_step": 0}
|
||||
assert spec.backend_kwargs() == {
|
||||
"sparsity": 0.8,
|
||||
"start_step": 0,
|
||||
"end_step": 0,
|
||||
"precision": "bf16",
|
||||
}
|
||||
|
||||
def test_block_sparse_skip_layers_selector_expanded(self):
|
||||
spec = AttentionSpec(
|
||||
@@ -104,6 +117,8 @@ class TestAttentionSpec:
|
||||
assert spec.backend_kwargs() == {
|
||||
"sparsity": 0.9,
|
||||
"start_step": 12,
|
||||
"end_step": 0,
|
||||
"precision": "bf16",
|
||||
"skip_layers": [0, 1, 2, 38],
|
||||
}
|
||||
|
||||
@@ -113,12 +128,44 @@ class TestAttentionSpec:
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"block_sparse",
|
||||
[{"sparsity": 1.5}, {"start_step": -1}],
|
||||
[{"sparsity": 1.5}, {"start_step": -1}, {"end_step": -1}, {"precision": "bogus"}],
|
||||
)
|
||||
def test_block_sparse_invalid_values(self, block_sparse):
|
||||
with pytest.raises(ValueError):
|
||||
AttentionSpec(backend="RAINFUSION_ATTN", block_sparse=block_sparse)
|
||||
|
||||
def test_block_sparse_end_step_precision_forwarded(self):
|
||||
spec = AttentionSpec(
|
||||
backend="RAINFUSION_ATTN",
|
||||
block_sparse={"sparsity": 0.9, "start_step": 5, "end_step": 3, "precision": "fp8"},
|
||||
)
|
||||
assert spec.backend_kwargs() == {
|
||||
"sparsity": 0.9,
|
||||
"start_step": 5,
|
||||
"end_step": 3,
|
||||
"precision": "fp8",
|
||||
}
|
||||
|
||||
def test_block_sparse_precision_enum_input_normalized(self):
|
||||
"""RainFusionPrecision enum members must normalize to their string value.
|
||||
|
||||
Regression: the str,Enum members passed the value check but str() later
|
||||
produced "RainFusionPrecision.MIX" instead of "mix" (StrEnum fixes it).
|
||||
"""
|
||||
from vllm_omni.diffusion.attention.backends.rainfusion_attn import RainFusionConfig
|
||||
from vllm_omni.diffusion.data import RainFusionPrecision
|
||||
|
||||
spec = AttentionSpec(
|
||||
backend="RAINFUSION_ATTN",
|
||||
block_sparse={"sparsity": 0.8, "precision": RainFusionPrecision.MIX},
|
||||
)
|
||||
kwargs = spec.backend_kwargs()
|
||||
assert kwargs["precision"] == "mix"
|
||||
# The config path is where str(value) used to mangle enum members.
|
||||
cfg = RainFusionConfig.from_backend_kwargs(kwargs)
|
||||
assert cfg.precision == "mix"
|
||||
assert type(cfg.precision) is str
|
||||
|
||||
|
||||
class TestAttentionConfig:
|
||||
def test_empty_config(self):
|
||||
|
||||
@@ -5,6 +5,7 @@ from contextlib import nullcontext
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch.nn.attention import SDPBackend
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
|
||||
@@ -12,6 +13,48 @@ from vllm_omni.diffusion.attention.backends.cudnn_attn import CuDNNAttentionImpl
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("query_sequence_length", "key_sequence_length", "expected_backend"),
|
||||
[
|
||||
(1, 1, SDPBackend.MATH),
|
||||
(8, 1, SDPBackend.MATH),
|
||||
(1, 8, SDPBackend.MATH),
|
||||
(8, 8, SDPBackend.CUDNN_ATTENTION),
|
||||
],
|
||||
)
|
||||
def test_cudnn_routes_single_token_attention_to_math(
|
||||
monkeypatch, query_sequence_length, key_sequence_length, expected_backend
|
||||
):
|
||||
selected_backends = []
|
||||
|
||||
def fake_sdpa_kernel(backends):
|
||||
selected_backends.append(backends)
|
||||
return nullcontext()
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.attention.backends.cudnn_attn.sdpa_kernel",
|
||||
fake_sdpa_kernel,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
torch.nn.functional,
|
||||
"scaled_dot_product_attention",
|
||||
lambda query, _key, _value, **_kwargs: query,
|
||||
)
|
||||
impl = CuDNNAttentionImpl(
|
||||
num_heads=2,
|
||||
head_size=4,
|
||||
softmax_scale=0.5,
|
||||
)
|
||||
query = torch.randn(1, query_sequence_length, 2, 4)
|
||||
key = torch.randn(1, key_sequence_length, 2, 4)
|
||||
value = torch.randn_like(key)
|
||||
|
||||
output = impl.forward_cuda(query, key, value)
|
||||
|
||||
assert output.shape == query.shape
|
||||
assert selected_backends == [[expected_backend]]
|
||||
|
||||
|
||||
def test_cudnn_slices_valid_kv_prefix_without_padding_mask(monkeypatch):
|
||||
observed = {}
|
||||
|
||||
|
||||
@@ -0,0 +1,203 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import sys
|
||||
import types
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.backends.fastvideo_vsa import (
|
||||
FastVideoVSAImpl,
|
||||
)
|
||||
from vllm_omni.diffusion.attention.backends.registry import (
|
||||
DiffusionAttentionBackendEnum,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_fastvideo_vsa_backend_is_registered():
|
||||
assert DiffusionAttentionBackendEnum.FASTVIDEO_VSA.get_path().endswith("fastvideo_vsa.FastVideoVSABackend")
|
||||
|
||||
|
||||
def test_fastvideo_vsa_tiles_3d_sequence_and_untiles(monkeypatch):
|
||||
calls: dict[str, Any] = {}
|
||||
fake_module = types.ModuleType("fastvideo_kernel")
|
||||
|
||||
def fake_video_sparse_attn_bshd(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
variable_block_sizes,
|
||||
q_variable_block_sizes,
|
||||
compress_attn_weight,
|
||||
topk,
|
||||
block_size,
|
||||
):
|
||||
calls["q_shape"] = tuple(q.shape)
|
||||
calls["vbs"] = variable_block_sizes.detach().cpu().tolist()
|
||||
calls["q_vbs"] = q_variable_block_sizes.detach().cpu().tolist()
|
||||
calls["compress_sum"] = float(compress_attn_weight.detach().cpu().sum())
|
||||
calls["compress_shape"] = tuple(compress_attn_weight.shape)
|
||||
calls["topk"] = topk
|
||||
calls["block_size"] = block_size
|
||||
return q + k + v
|
||||
|
||||
setattr(fake_module, "video_sparse_attn_bshd", fake_video_sparse_attn_bshd)
|
||||
monkeypatch.setitem(sys.modules, "fastvideo_kernel", fake_module)
|
||||
|
||||
impl = FastVideoVSAImpl(
|
||||
num_heads=2,
|
||||
head_size=8,
|
||||
softmax_scale=8**-0.5,
|
||||
causal=False,
|
||||
backend_kwargs={
|
||||
"topk": 1,
|
||||
"block_size": (4, 8, 8),
|
||||
"min_seq_len": 1,
|
||||
"disable_when_sp_active": False,
|
||||
},
|
||||
)
|
||||
query = torch.randn(1, 300, 2, 8)
|
||||
key = torch.randn_like(query)
|
||||
value = torch.randn_like(query)
|
||||
attn_metadata = AttentionMetadata(extra={"vsa_dit_seq_shape": (3, 10, 10)})
|
||||
monkeypatch.setattr(impl, "_fallback_reason", lambda *args, **kwargs: None)
|
||||
|
||||
output = impl.forward_cuda(query, key, value, attn_metadata)
|
||||
|
||||
assert output.shape == query.shape
|
||||
assert calls["q_shape"] == (1, 1024, 2, 8)
|
||||
assert calls["vbs"] == [192, 48, 48, 12]
|
||||
assert calls["q_vbs"] == [192, 48, 48, 12]
|
||||
assert calls["compress_shape"] == (1, 1024, 2, 8)
|
||||
assert calls["compress_sum"] == 0.0
|
||||
assert calls["topk"] == 1
|
||||
assert calls["block_size"] == (4, 8, 8)
|
||||
|
||||
|
||||
def test_fastvideo_vsa_uses_learned_gate_when_provided(monkeypatch):
|
||||
calls: dict[str, Any] = {}
|
||||
fake_module = types.ModuleType("fastvideo_kernel")
|
||||
|
||||
def fake_video_sparse_attn_bshd(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
variable_block_sizes,
|
||||
q_variable_block_sizes,
|
||||
compress_attn_weight,
|
||||
topk,
|
||||
block_size,
|
||||
):
|
||||
calls["compress_sum"] = float(compress_attn_weight.detach().cpu().sum())
|
||||
calls["compress_shape"] = tuple(compress_attn_weight.shape)
|
||||
return q + k + v
|
||||
|
||||
setattr(fake_module, "video_sparse_attn_bshd", fake_video_sparse_attn_bshd)
|
||||
monkeypatch.setitem(sys.modules, "fastvideo_kernel", fake_module)
|
||||
|
||||
impl = FastVideoVSAImpl(
|
||||
num_heads=2,
|
||||
head_size=8,
|
||||
softmax_scale=8**-0.5,
|
||||
causal=False,
|
||||
backend_kwargs={
|
||||
"topk": 1,
|
||||
"block_size": (4, 8, 8),
|
||||
"min_seq_len": 1,
|
||||
"disable_when_sp_active": False,
|
||||
},
|
||||
)
|
||||
query = torch.randn(1, 300, 2, 8)
|
||||
gate = torch.ones_like(query)
|
||||
metadata = AttentionMetadata(extra={"vsa_dit_seq_shape": (3, 10, 10), "gate_compress": gate})
|
||||
monkeypatch.setattr(impl, "_fallback_reason", lambda *args, **kwargs: None)
|
||||
|
||||
output = impl.forward_cuda(query, query, query, metadata)
|
||||
|
||||
assert output.shape == query.shape
|
||||
assert calls["compress_shape"] == (1, 1024, 2, 8)
|
||||
# Only the 300 real tokens carry ones; padded gate slots stay zero.
|
||||
assert calls["compress_sum"] == gate.numel()
|
||||
|
||||
|
||||
def test_fastvideo_vsa_falls_back_without_dit_shape(monkeypatch):
|
||||
calls: dict[str, Any] = {}
|
||||
|
||||
impl = FastVideoVSAImpl(
|
||||
num_heads=2,
|
||||
head_size=8,
|
||||
softmax_scale=8**-0.5,
|
||||
causal=False,
|
||||
backend_kwargs={
|
||||
"topk": 1,
|
||||
"block_size": (4, 8, 8),
|
||||
"min_seq_len": 1,
|
||||
"disable_when_sp_active": False,
|
||||
},
|
||||
)
|
||||
|
||||
def fake_fallback(query, key, value, attn_metadata, reason):
|
||||
calls["reason"] = reason
|
||||
return torch.zeros_like(query)
|
||||
|
||||
monkeypatch.setattr(impl, "_fallback", fake_fallback)
|
||||
|
||||
query = torch.randn(1, 512, 2, 8)
|
||||
output = impl.forward_cuda(query, query, query, AttentionMetadata())
|
||||
|
||||
assert output.shape == query.shape
|
||||
assert calls["reason"] == "vsa_dit_seq_shape metadata is required"
|
||||
|
||||
|
||||
def test_fastvideo_vsa_falls_back_for_mask(monkeypatch):
|
||||
calls: dict[str, Any] = {}
|
||||
|
||||
impl = FastVideoVSAImpl(
|
||||
num_heads=2,
|
||||
head_size=8,
|
||||
softmax_scale=8**-0.5,
|
||||
causal=False,
|
||||
backend_kwargs={"min_seq_len": 1},
|
||||
)
|
||||
|
||||
def fake_fallback(query, key, value, attn_metadata, reason):
|
||||
calls["reason"] = reason
|
||||
return torch.zeros_like(query)
|
||||
|
||||
monkeypatch.setattr(impl, "_fallback", fake_fallback)
|
||||
|
||||
query = torch.randn(1, 512, 2, 8)
|
||||
mask = torch.ones(1, 512, dtype=torch.bool)
|
||||
output = impl.forward_cuda(query, query, query, AttentionMetadata(attn_mask=mask))
|
||||
|
||||
assert output.shape == query.shape
|
||||
assert calls["reason"]
|
||||
|
||||
|
||||
def test_fastvideo_vsa_allows_topk_equal_to_num_blocks():
|
||||
query = torch.randn(1, 512, 2, 8)
|
||||
metadata = AttentionMetadata(extra={"vsa_dit_seq_shape": (4, 8, 16)})
|
||||
|
||||
all_blocks = FastVideoVSAImpl(
|
||||
num_heads=2,
|
||||
head_size=8,
|
||||
softmax_scale=8**-0.5,
|
||||
causal=False,
|
||||
backend_kwargs={"topk": 2, "block_size": (4, 8, 8), "min_seq_len": 1},
|
||||
)
|
||||
too_many = FastVideoVSAImpl(
|
||||
num_heads=2,
|
||||
head_size=8,
|
||||
softmax_scale=8**-0.5,
|
||||
causal=False,
|
||||
backend_kwargs={"topk": 3, "block_size": (4, 8, 8), "min_seq_len": 1},
|
||||
)
|
||||
|
||||
# CPU/float32 is rejected later, but k=N itself must not trigger fallback.
|
||||
assert all_blocks._fallback_reason(query, query, query, metadata) == "dtype torch.float32 is not supported"
|
||||
assert too_many._fallback_reason(query, query, query, metadata) == "topk 3 > num_blocks 2"
|
||||
@@ -265,10 +265,15 @@ def test_fa_vs_sdpa():
|
||||
print(f"Max absolute difference (valid region): {max_diff:.6f}")
|
||||
print(f"Mean absolute difference (valid region): {mean_diff:.6f}")
|
||||
|
||||
# Assert that outputs are close
|
||||
# Using higher tolerance for bfloat16 and different implementations
|
||||
assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
|
||||
assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
|
||||
# AITER's ROCm BF16 kernel is deterministic, but differs from ROCm SDPA by
|
||||
# about 1.3% relative error. An absolute-only bound misclassifies that
|
||||
# expected rounding difference for larger output values (for example,
|
||||
# -2.65625 vs -2.6875). Keep the tighter historical bounds elsewhere.
|
||||
if current_omni_platform.is_rocm():
|
||||
torch.testing.assert_close(output_fa_valid, output_sdpa_valid, rtol=1e-2, atol=1.5e-2)
|
||||
else:
|
||||
assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
|
||||
assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
|
||||
|
||||
print("✓ Case 2 PASSED: FA and SDPA outputs are very close!")
|
||||
|
||||
@@ -303,6 +308,91 @@ def test_flash_attn_func_preferred_over_varlen():
|
||||
print("✓ flash_attn_func forward works correctly!")
|
||||
|
||||
|
||||
@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires CUDA or XPU")
|
||||
@pytest.mark.parametrize("k_len", [40, 256])
|
||||
def test_cross_attn_key_padding_vs_sdpa(k_len):
|
||||
"""
|
||||
Case 3: Cross-attention with a key-padding mask.
|
||||
|
||||
The mask covers only the key sequence; every query position must attend
|
||||
to the valid keys. FA and SDPA outputs should be very close. k_len=256
|
||||
covers equal Q/K lengths, which must not be mistaken for self-attention,
|
||||
and valid key lengths differ per batch element.
|
||||
"""
|
||||
device = torch.device(current_omni_platform.device_type)
|
||||
dtype = torch.bfloat16
|
||||
|
||||
batch_size = 2
|
||||
q_len = 256
|
||||
valid_lens = (17, 29)
|
||||
num_heads = 8
|
||||
head_dim = 64
|
||||
|
||||
fa_impl = FlashAttentionImpl(
|
||||
num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False, role="cross"
|
||||
)
|
||||
sdpa_impl = SDPAImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False)
|
||||
|
||||
torch.manual_seed(7)
|
||||
query = torch.randn(batch_size, q_len, num_heads, head_dim, device=device, dtype=dtype)
|
||||
key = torch.randn(batch_size, k_len, num_heads, head_dim, device=device, dtype=dtype)
|
||||
value = torch.randn(batch_size, k_len, num_heads, head_dim, device=device, dtype=dtype)
|
||||
|
||||
attn_mask = torch.zeros(batch_size, k_len, dtype=torch.bool, device=device)
|
||||
for i, valid_len in enumerate(valid_lens):
|
||||
attn_mask[i, :valid_len] = True
|
||||
attn_metadata = AttentionMetadata(attn_mask=attn_mask)
|
||||
|
||||
output_fa = fa_impl.forward(query=query, key=key, value=value, attn_metadata=attn_metadata)
|
||||
output_sdpa = sdpa_impl.forward(query=query, key=key, value=value, attn_metadata=attn_metadata)
|
||||
|
||||
max_diff = torch.max(torch.abs(output_fa - output_sdpa)).item()
|
||||
mean_diff = torch.mean(torch.abs(output_fa - output_sdpa)).item()
|
||||
|
||||
print("\n=== Case 3: Cross-Attention Key-Padding FA vs SDPA ===")
|
||||
print(f"Max absolute difference: {max_diff:.6f}")
|
||||
print(f"Mean absolute difference: {mean_diff:.6f}")
|
||||
|
||||
assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01"
|
||||
assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001"
|
||||
|
||||
print("✓ Case 3 PASSED: FA and SDPA cross-attention outputs are very close!")
|
||||
|
||||
|
||||
def test_varlen_masked_routing_by_role(monkeypatch):
|
||||
"""The unpad route is picked by role, not Q/K length equality; runs without a GPU."""
|
||||
calls = []
|
||||
|
||||
def fake_varlen_func(q, k, v, **kwargs):
|
||||
calls.append((q.shape[0], kwargs["cu_seqlens_q"], kwargs["cu_seqlens_k"]))
|
||||
return q
|
||||
|
||||
monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True)
|
||||
monkeypatch.setattr(fa, "flash_attn_varlen_func", fake_varlen_func)
|
||||
|
||||
batch_size, seq_len, num_heads, head_dim = 2, 4, 2, 8
|
||||
query = torch.randn(batch_size, seq_len, num_heads, head_dim)
|
||||
metadata = AttentionMetadata(attn_mask=torch.tensor([[True, True, False, False], [True, True, True, False]]))
|
||||
|
||||
cross_impl = FlashAttentionImpl(
|
||||
num_heads=num_heads, head_size=head_dim, softmax_scale=0.5, causal=False, role="cross"
|
||||
)
|
||||
output = cross_impl.forward_cuda(query, query, query, metadata)
|
||||
assert output.shape == query.shape
|
||||
num_q_rows, cu_seqlens_q, cu_seqlens_k = calls[-1]
|
||||
assert num_q_rows == batch_size * seq_len
|
||||
assert cu_seqlens_q.tolist() == [0, 4, 8]
|
||||
assert cu_seqlens_k.tolist() == [0, 2, 5]
|
||||
|
||||
self_impl = FlashAttentionImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=0.5, causal=False)
|
||||
output = self_impl.forward_cuda(query, query, query, metadata)
|
||||
assert output.shape == query.shape
|
||||
num_q_rows, cu_seqlens_q, cu_seqlens_k = calls[-1]
|
||||
assert num_q_rows == 5
|
||||
assert cu_seqlens_q.tolist() == [0, 2, 5]
|
||||
assert cu_seqlens_k.tolist() == [0, 2, 5]
|
||||
|
||||
|
||||
def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch):
|
||||
calls = []
|
||||
|
||||
|
||||
@@ -23,8 +23,10 @@ import vllm_omni.diffusion.attention.backends.utils.piecewise_attn as piecewise_
|
||||
from vllm_omni.diffusion.attention.backends.abstract import QueryRange
|
||||
from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
|
||||
Segment,
|
||||
build_paged_piecewise_plan,
|
||||
build_segments,
|
||||
piecewise_attn,
|
||||
run_paged_piecewise_plan,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
@@ -84,6 +86,72 @@ def test_build_segments_full_span_retains_global_kv_end():
|
||||
assert segments == [Segment(q_start=4, q_end=6, kv_end=7, mode="full")]
|
||||
|
||||
|
||||
def test_paged_piecewise_runner_uses_packed_indices_and_kv_endpoints():
|
||||
full_attn_spans = [[(2, 5)], [(4, 7)]]
|
||||
plan = build_paged_piecewise_plan(
|
||||
full_attn_spans,
|
||||
query_offsets=[0, 2],
|
||||
query_lens=[6, 6],
|
||||
seq_lens=[6, 8],
|
||||
device=DEVICE,
|
||||
)
|
||||
query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4)
|
||||
key = query + 100
|
||||
value = query + 200
|
||||
calls = []
|
||||
|
||||
def run_segment(segment_query, segment_key, segment_value, segment_batch):
|
||||
calls.append(
|
||||
(
|
||||
[segment.mode for segment in segment_batch.row_segments],
|
||||
[segment.kv_end for segment in segment_batch.row_segments],
|
||||
segment_batch.query_indices.tolist(),
|
||||
)
|
||||
)
|
||||
torch.testing.assert_close(segment_key, segment_query + 100)
|
||||
torch.testing.assert_close(segment_value, segment_query + 200)
|
||||
return segment_query
|
||||
|
||||
output = run_paged_piecewise_plan(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
plan,
|
||||
plan.segments,
|
||||
run_segment,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(output, query)
|
||||
assert calls == [
|
||||
(["causal", "causal"], [2, 4], [0, 1, 6, 7]),
|
||||
(["full", "full"], [5, 7], [2, 3, 4, 8, 9, 10]),
|
||||
(["causal", "causal"], [6, 8], [5, 11]),
|
||||
]
|
||||
|
||||
|
||||
def test_paged_piecewise_runner_uses_runner_output_shape():
|
||||
plan = build_paged_piecewise_plan(
|
||||
[[(2, 5)]],
|
||||
query_offsets=[0],
|
||||
query_lens=[6],
|
||||
seq_lens=[6],
|
||||
device=DEVICE,
|
||||
)
|
||||
query = torch.arange(6 * 2 * 4, dtype=torch.float32).reshape(6, 2, 4)
|
||||
|
||||
output = run_paged_piecewise_plan(
|
||||
query,
|
||||
query,
|
||||
query,
|
||||
plan,
|
||||
plan.segments,
|
||||
lambda segment_query, *_args: segment_query[..., :3],
|
||||
)
|
||||
|
||||
assert output.shape == (6, 2, 3)
|
||||
torch.testing.assert_close(output, query[..., :3])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("global_spans", SPAN_CASES)
|
||||
@pytest.mark.parametrize("q_range", Q_RANGE_CASES)
|
||||
@pytest.mark.parametrize("batch_size", BATCH_CASES)
|
||||
|
||||
@@ -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
|
||||
|
||||
"""Geometry contract RAINFUSION_ATTN enforces before handing a forward to rf_v2.
|
||||
|
||||
@@ -11,18 +11,25 @@ to the always-kept prefix segment before generating the block mask.
|
||||
"""
|
||||
|
||||
import dataclasses
|
||||
import sys
|
||||
import types
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout
|
||||
from vllm_omni.diffusion.attention.backends import rainfusion_attn
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout, VideoTokenSpan
|
||||
from vllm_omni.diffusion.attention.backends.rainfusion_attn import (
|
||||
_BLOCK_SIZE,
|
||||
RainFusionAttentionBackend,
|
||||
RainFusionAttentionImpl,
|
||||
RainFusionPlan,
|
||||
)
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.npu]
|
||||
|
||||
PREFIX_ROWS = 710 # 14 text rows + 696 audio rows
|
||||
ALIGNED_GRID = (62, 24, 40) # 1280x768 -> 59520 video rows, 465 blocks
|
||||
MISALIGNED_GRID = (62, 24, 42) # 1344x768 -> 62496 video rows, 488 blocks + 32 rows
|
||||
@@ -97,6 +104,62 @@ def test_video_segment_must_be_the_tail_of_packed_document_zero():
|
||||
assert make_impl()._resolve_plan(metadata) is None
|
||||
|
||||
|
||||
def test_ref2va_multi_video_spans_resolve():
|
||||
layout = VideoTokenLayout(
|
||||
used_len=12000,
|
||||
video_spans=(
|
||||
VideoTokenSpan(start=128, latent_grid=(4, 16, 64), role="reference"),
|
||||
VideoTokenSpan(start=5000, latent_grid=(4, 16, 64), role="target"),
|
||||
),
|
||||
)
|
||||
plan = make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 12000}, video_layout=layout))
|
||||
|
||||
assert plan is not None
|
||||
assert plan.used_len == 12000
|
||||
assert plan.video_spans == [
|
||||
{"start": 128, "latent_shape": [4, 16, 64]},
|
||||
{"start": 5000, "latent_shape": [4, 16, 64]},
|
||||
]
|
||||
|
||||
|
||||
def test_invalid_ref2va_video_spans_fall_back_to_dense():
|
||||
layout = VideoTokenLayout(
|
||||
used_len=12000,
|
||||
video_spans=(
|
||||
VideoTokenSpan(start=128, latent_grid=(4, 16, 64), role="reference"),
|
||||
VideoTokenSpan(start=4000, latent_grid=(4, 16, 64), role="target"),
|
||||
),
|
||||
)
|
||||
assert make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 12000}, video_layout=layout)) is None
|
||||
|
||||
|
||||
def test_validate_available_rejects_legacy_mindiesd(monkeypatch):
|
||||
mindiesd = types.ModuleType("mindiesd")
|
||||
|
||||
def sparse_attention(query, key, value, **kwargs):
|
||||
return query
|
||||
|
||||
mindiesd.sparse_attention = sparse_attention
|
||||
monkeypatch.setitem(sys.modules, "mindiesd", mindiesd)
|
||||
monkeypatch.setattr("importlib.util.find_spec", lambda _: object())
|
||||
|
||||
with pytest.raises(ValueError, match="video_spans"):
|
||||
RainFusionAttentionBackend.validate_available()
|
||||
|
||||
|
||||
def test_validate_available_accepts_new_mindiesd(monkeypatch):
|
||||
mindiesd = types.ModuleType("mindiesd")
|
||||
|
||||
def sparse_attention(query, key, value, *, video_spans=None, **kwargs):
|
||||
return query
|
||||
|
||||
mindiesd.sparse_attention = sparse_attention
|
||||
monkeypatch.setitem(sys.modules, "mindiesd", mindiesd)
|
||||
monkeypatch.setattr("importlib.util.find_spec", lambda _: object())
|
||||
|
||||
RainFusionAttentionBackend.validate_available()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("grid", [(4, 24, 40), (1, 24, 40)])
|
||||
def test_short_video_stays_dense(grid):
|
||||
assert make_impl()._resolve_plan(make_metadata(grid)) is None
|
||||
@@ -169,7 +232,9 @@ def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd()
|
||||
heads, head_dim = 4, 128
|
||||
q, k, v = (torch.randn(1, used, heads, head_dim, dtype=torch.bfloat16, device="npu") for _ in range(3))
|
||||
|
||||
impl = make_impl()
|
||||
# Explicit bf16: the dense comparison exercises the plain sparse kernel,
|
||||
# and older MindIE-SD releases do not accept precision= (gate would raise).
|
||||
impl = make_impl(precision="bf16")
|
||||
impl.rainfusion = dataclasses.replace(impl.rainfusion, sparsity=0.0)
|
||||
plan = RainFusionPlan(prefix_len=prefix_len, used_len=used, latent_shape=list(grid))
|
||||
out = impl._forward_sparse_npu(q, k, v, plan)
|
||||
@@ -182,3 +247,90 @@ def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd()
|
||||
).transpose(1, 2)
|
||||
error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
|
||||
assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"
|
||||
|
||||
|
||||
# end_step tail fallback: the last ``end_step`` denoise steps must stay dense.
|
||||
|
||||
|
||||
def test_end_step_tail_window_stays_dense():
|
||||
"""step_idx inside the last end_step window must resolve no sparse plan."""
|
||||
impl = make_impl(end_step=3)
|
||||
fc = mock.Mock()
|
||||
fc.denoise_step_idx = 47
|
||||
fc.total_denoise_steps = 50
|
||||
with (
|
||||
mock.patch.object(rainfusion_attn, "is_forward_context_available", return_value=True),
|
||||
mock.patch.object(rainfusion_attn, "get_forward_context", return_value=fc),
|
||||
):
|
||||
assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is None
|
||||
|
||||
|
||||
def test_end_step_before_tail_window_still_sparse():
|
||||
"""step_idx before the tail window must still resolve a sparse plan."""
|
||||
impl = make_impl(end_step=3)
|
||||
fc = mock.Mock()
|
||||
fc.denoise_step_idx = 46
|
||||
fc.total_denoise_steps = 50
|
||||
with (
|
||||
mock.patch.object(rainfusion_attn, "is_forward_context_available", return_value=True),
|
||||
mock.patch.object(rainfusion_attn, "get_forward_context", return_value=fc),
|
||||
):
|
||||
assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is not None
|
||||
|
||||
|
||||
def test_end_step_zero_never_triggers_tail_fallback():
|
||||
"""end_step=0 must not fall back even on the final denoise step."""
|
||||
impl = make_impl(end_step=0)
|
||||
fc = mock.Mock()
|
||||
fc.denoise_step_idx = 49
|
||||
fc.total_denoise_steps = 50
|
||||
with (
|
||||
mock.patch.object(rainfusion_attn, "is_forward_context_available", return_value=True),
|
||||
mock.patch.object(rainfusion_attn, "get_forward_context", return_value=fc),
|
||||
):
|
||||
assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is not None
|
||||
|
||||
|
||||
# precision capability gate: a mindiesd without sparse_attention(precision=)
|
||||
# must fail fast instead of silently running the BF16 path.
|
||||
|
||||
|
||||
def _fake_mindiesd_module():
|
||||
"""Install a minimal stand-in so the import inside _forward_sparse_npu succeeds."""
|
||||
import types
|
||||
|
||||
fake = types.ModuleType("mindiesd")
|
||||
fake.sparse_attention = lambda *args, **kwargs: None
|
||||
return fake
|
||||
|
||||
|
||||
def test_precision_non_bf16_requires_mindiesd_support():
|
||||
"""precision != bf16 against a mindiesd lacking the kwarg must raise RuntimeError."""
|
||||
import sys
|
||||
|
||||
impl = make_impl(precision="mix")
|
||||
sys.modules["mindiesd"] = _fake_mindiesd_module()
|
||||
try:
|
||||
with mock.patch.object(rainfusion_attn, "_mindiesd_supports_precision", return_value=False):
|
||||
with pytest.raises(RuntimeError, match="requires MindIE-SD"):
|
||||
impl._forward_sparse_npu(None, None, None, None)
|
||||
finally:
|
||||
sys.modules.pop("mindiesd", None)
|
||||
|
||||
|
||||
def test_precision_non_bf16_passes_gate_when_supported():
|
||||
"""precision != bf16 with a capable mindiesd must not raise from the gate."""
|
||||
import sys
|
||||
|
||||
impl = make_impl(precision="mix")
|
||||
sys.modules["mindiesd"] = _fake_mindiesd_module()
|
||||
try:
|
||||
with mock.patch.object(rainfusion_attn, "_mindiesd_supports_precision", return_value=True):
|
||||
# q/k/v shapes: [B, S, N, D]; plan geometry must match S.
|
||||
q = torch.randn(1, 8, 4, 128)
|
||||
plan = RainFusionPlan(prefix_len=0, used_len=8, latent_shape=[2, 2, 2])
|
||||
# The gate must pass; the fake mindiesd returns None so no crash.
|
||||
out = impl._forward_sparse_npu(q, q, q, plan)
|
||||
assert out is None
|
||||
finally:
|
||||
sys.modules.pop("mindiesd", None)
|
||||
@@ -1,14 +1,18 @@
|
||||
# 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 unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from vllm_omni.diffusion.attention.backends import trtllm_attn as tg
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.backends.abstract import (
|
||||
AttentionMetadata,
|
||||
PackedPaddingMetadata,
|
||||
)
|
||||
from vllm_omni.diffusion.attention.backends.trtllm_attn import TrtllmAttentionImpl
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda]
|
||||
@@ -35,6 +39,15 @@ def _impl(**bk):
|
||||
return TrtllmAttentionImpl(8, 128, 1.0 / math.sqrt(128), causal=False, num_kv_heads=8, backend_kwargs=bk)
|
||||
|
||||
|
||||
def _packed_padding(cu_seqlens_q, cu_seqlens_k, q_length, kv_length):
|
||||
return PackedPaddingMetadata(
|
||||
q_length=q_length,
|
||||
kv_length=kv_length,
|
||||
cu_seqlens_q=cu_seqlens_q[:2],
|
||||
cu_seqlens_k=cu_seqlens_k[:2],
|
||||
)
|
||||
|
||||
|
||||
def test_skip_config_pure_resolution():
|
||||
from vllm_omni.diffusion.attention.backends.trtllm_attn import SkipSoftmaxConfig
|
||||
|
||||
@@ -75,7 +88,7 @@ def test_quant_quantize_requires_flashinfer_routine(monkeypatch):
|
||||
def test_quant_quantize_calls_routine_and_shapes_sfs():
|
||||
from vllm_omni.diffusion.attention.backends.trtllm_attn import QuantConfig
|
||||
|
||||
captured = {}
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_quantize(q, k, v, q_block_size, k_block_size, qk_quant_dtype):
|
||||
captured.update(q_block_size=q_block_size, k_block_size=k_block_size, qk_quant_dtype=qk_quant_dtype)
|
||||
@@ -209,62 +222,80 @@ def test_masked_layer_raises():
|
||||
_impl().forward_cuda(q, k, v, _Meta())
|
||||
|
||||
|
||||
def test_packed_metadata_trims_padding_mask(monkeypatch):
|
||||
b, s, h, d = 1, 8, 8, 128
|
||||
q, k, v = (torch.zeros(b, s, h, d) for _ in range(3))
|
||||
mask = torch.tensor([[True, True, True, True, True, True, False, False]])
|
||||
cu_seqlens = torch.tensor([0, 6, 8], dtype=torch.int32)
|
||||
def test_generic_packed_metadata_preserves_ragged_batch(monkeypatch):
|
||||
b, h, d = 1, 8, 128
|
||||
q = torch.zeros(b, 5, h, d)
|
||||
k, v = (torch.zeros(b, 7, h, d) for _ in range(2))
|
||||
metadata = AttentionMetadata(
|
||||
attn_mask=mask,
|
||||
extra={
|
||||
"cu_seqlens_q": cu_seqlens,
|
||||
"cu_seqlens_k": cu_seqlens,
|
||||
"max_seqlen_q": 6,
|
||||
"max_seqlen_k": 6,
|
||||
"cu_seqlens_q": torch.tensor([0, 2, 5], dtype=torch.int32),
|
||||
"cu_seqlens_k": torch.tensor([0, 3, 7], dtype=torch.int32),
|
||||
"max_seqlen_q": 3,
|
||||
"max_seqlen_k": 4,
|
||||
}
|
||||
)
|
||||
fake_attention = Mock(return_value=torch.zeros_like(q.reshape(5, h, d)))
|
||||
|
||||
monkeypatch.setattr(tg, "HAS_FLASHINFER", True)
|
||||
monkeypatch.setattr(tg, "trtllm_ragged_attention_deepseek", fake_attention)
|
||||
monkeypatch.setattr(TrtllmAttentionImpl, "_get_workspace", classmethod(lambda cls, device: torch.empty(0)))
|
||||
|
||||
_impl().forward_cuda(q, k, v, metadata)
|
||||
captured = fake_attention.call_args.kwargs
|
||||
|
||||
assert captured["query"].shape[0] == 5
|
||||
assert captured["key"].shape[0] == 7
|
||||
assert captured["value"].shape[0] == 7
|
||||
assert captured["batch_size"] == 2
|
||||
assert captured["seq_lens"].tolist() == [3, 4]
|
||||
assert tg.TrtllmAttentionBackend.supports_multi_doc_packed_varlen()
|
||||
|
||||
|
||||
def test_packed_padding_metadata_trims_without_mask(monkeypatch):
|
||||
q_len, kv_len = 8, 9
|
||||
valid_q_tokens, valid_kv_tokens = 5, 6
|
||||
b, h, d = 1, 8, 128
|
||||
q = torch.zeros(b, q_len, h, d)
|
||||
k, v = (torch.zeros(b, kv_len, h, d) for _ in range(2))
|
||||
cu_seqlens_q = torch.tensor([0, valid_q_tokens, q_len], dtype=torch.int32)
|
||||
cu_seqlens_k = torch.tensor([0, valid_kv_tokens, kv_len], dtype=torch.int32)
|
||||
metadata = AttentionMetadata(
|
||||
packed_padding=_packed_padding(
|
||||
cu_seqlens_q,
|
||||
cu_seqlens_k,
|
||||
valid_q_tokens,
|
||||
valid_kv_tokens,
|
||||
),
|
||||
extra={
|
||||
"cu_seqlens_q": cu_seqlens_q,
|
||||
"cu_seqlens_k": cu_seqlens_k,
|
||||
"max_seqlen_q": valid_q_tokens,
|
||||
"max_seqlen_k": valid_kv_tokens,
|
||||
"valid_kv_length": valid_kv_tokens,
|
||||
},
|
||||
)
|
||||
captured = {}
|
||||
|
||||
def fake_attention(**kwargs):
|
||||
captured.update(kwargs)
|
||||
return torch.zeros_like(kwargs["query"])
|
||||
fake_attention = Mock(return_value=torch.ones(valid_q_tokens, h, d))
|
||||
|
||||
monkeypatch.setattr(tg, "HAS_FLASHINFER", True)
|
||||
monkeypatch.setattr(tg, "trtllm_ragged_attention_deepseek", fake_attention)
|
||||
monkeypatch.setattr(TrtllmAttentionImpl, "_get_workspace", classmethod(lambda cls, device: torch.empty(0)))
|
||||
|
||||
out = _impl().forward_cuda(q, k, v, metadata)
|
||||
captured = fake_attention.call_args.kwargs
|
||||
|
||||
assert tg.TrtllmAttentionBackend.supports_packed_mask_free()
|
||||
assert out.shape == q.shape
|
||||
assert captured["query"].shape[0] == 6
|
||||
assert captured["key"].shape[0] == 6
|
||||
assert captured["value"].shape[0] == 6
|
||||
assert captured["query"].shape[0] == valid_q_tokens
|
||||
assert captured["key"].shape[0] == valid_kv_tokens
|
||||
assert captured["value"].shape[0] == valid_kv_tokens
|
||||
assert captured["batch_size"] == 1
|
||||
assert captured["seq_lens"].tolist() == [6]
|
||||
assert captured["cum_seq_lens_q"].tolist() == [0, 6]
|
||||
assert captured["cum_seq_lens_kv"].tolist() == [0, 6]
|
||||
assert captured["max_q_len"] == 6
|
||||
assert captured["max_kv_len"] == 6
|
||||
assert torch.count_nonzero(out[:, 6:]) == 0
|
||||
|
||||
|
||||
def test_packed_metadata_rejects_non_prefix_mask(monkeypatch):
|
||||
b, s, h, d = 1, 8, 8, 128
|
||||
q, k, v = (torch.zeros(b, s, h, d) for _ in range(3))
|
||||
cu_seqlens = torch.tensor([0, 6, 8], dtype=torch.int32)
|
||||
metadata = AttentionMetadata(
|
||||
attn_mask=torch.tensor([[True, True, False, True, True, True, False, False]]),
|
||||
extra={
|
||||
"cu_seqlens_q": cu_seqlens,
|
||||
"cu_seqlens_k": cu_seqlens,
|
||||
"max_seqlen_q": 6,
|
||||
"max_seqlen_k": 6,
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(tg, "HAS_FLASHINFER", True)
|
||||
|
||||
with pytest.raises(ValueError, match="prefix-valid"):
|
||||
_impl().forward_cuda(q, k, v, metadata)
|
||||
assert captured["seq_lens"].tolist() == [valid_kv_tokens]
|
||||
assert captured["cum_seq_lens_q"].tolist() == [0, valid_q_tokens]
|
||||
assert captured["cum_seq_lens_kv"].tolist() == [0, valid_kv_tokens]
|
||||
assert captured["max_q_len"] == valid_q_tokens
|
||||
assert captured["max_kv_len"] == valid_kv_tokens
|
||||
assert torch.count_nonzero(out[:, :valid_q_tokens] != 1) == 0
|
||||
assert torch.count_nonzero(out[:, valid_q_tokens:]) == 0
|
||||
|
||||
|
||||
def test_packed_metadata_must_be_complete():
|
||||
@@ -393,12 +424,13 @@ def test_bf16_packed_padding_matches_sdpa():
|
||||
q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) for _ in range(3))
|
||||
cu_seqlens = torch.tensor([0, used, s], dtype=torch.int32, device="cuda")
|
||||
metadata = AttentionMetadata(
|
||||
attn_mask=torch.arange(s, device="cuda")[None] < used,
|
||||
packed_padding=_packed_padding(cu_seqlens, cu_seqlens, used, used),
|
||||
extra={
|
||||
"cu_seqlens_q": cu_seqlens,
|
||||
"cu_seqlens_k": cu_seqlens,
|
||||
"max_seqlen_q": used,
|
||||
"max_seqlen_k": used,
|
||||
"valid_kv_length": used,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -443,12 +475,13 @@ def test_sage_packed_non_aligned_length_matches_sdpa():
|
||||
q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) for _ in range(3))
|
||||
cu_seqlens = torch.tensor([0, used, s], dtype=torch.int32, device="cuda")
|
||||
metadata = AttentionMetadata(
|
||||
attn_mask=torch.arange(s, device="cuda")[None] < used,
|
||||
packed_padding=_packed_padding(cu_seqlens, cu_seqlens, used, used),
|
||||
extra={
|
||||
"cu_seqlens_q": cu_seqlens,
|
||||
"cu_seqlens_k": cu_seqlens,
|
||||
"max_seqlen_q": used,
|
||||
"max_seqlen_k": used,
|
||||
"valid_kv_length": used,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@@ -1,10 +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
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import socket
|
||||
import tempfile
|
||||
from types import SimpleNamespace
|
||||
|
||||
@@ -12,6 +11,7 @@ import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.helpers.runtime import get_distributed_init_method
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention
|
||||
from vllm_omni.diffusion.config import set_current_diffusion_config
|
||||
@@ -27,20 +27,6 @@ from vllm_omni.platforms import current_omni_platform
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _find_free_port() -> int:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
return int(s.getsockname()[1])
|
||||
|
||||
|
||||
def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None:
|
||||
os.environ["RANK"] = str(rank)
|
||||
os.environ["LOCAL_RANK"] = str(rank)
|
||||
os.environ["WORLD_SIZE"] = str(world_size)
|
||||
os.environ["MASTER_ADDR"] = "127.0.0.1"
|
||||
os.environ["MASTER_PORT"] = str(master_port)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_advanced_uaa_hybrid_rejects_padded_mqa_heads(monkeypatch) -> None:
|
||||
@@ -71,7 +57,7 @@ def test_advanced_uaa_hybrid_rejects_padded_mqa_heads(monkeypatch) -> None:
|
||||
def _run_attention_case(
|
||||
local_rank: int,
|
||||
world_size: int,
|
||||
master_port: int,
|
||||
init_method: str,
|
||||
input_file: str,
|
||||
output_file: str,
|
||||
num_heads: int,
|
||||
@@ -85,10 +71,9 @@ def _run_attention_case(
|
||||
device = torch.device(f"{current_omni_platform.device_type}:{local_rank}")
|
||||
current_omni_platform.set_device(device)
|
||||
|
||||
_set_dist_env(rank=local_rank, world_size=world_size, master_port=master_port)
|
||||
os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA"
|
||||
|
||||
init_distributed_environment(world_size=world_size, rank=local_rank)
|
||||
init_distributed_environment(world_size=world_size, rank=local_rank, distributed_init_method=init_method)
|
||||
initialize_model_parallel(
|
||||
data_parallel_size=1,
|
||||
cfg_parallel_size=1,
|
||||
@@ -202,10 +187,8 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head
|
||||
batch_size = 2
|
||||
head_size = 8
|
||||
|
||||
base_port = _find_free_port()
|
||||
sp_port = _find_free_port()
|
||||
while sp_port == base_port:
|
||||
sp_port = _find_free_port()
|
||||
base_init_method = get_distributed_init_method()
|
||||
sp_init_method = get_distributed_init_method()
|
||||
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in:
|
||||
input_file = f_in.name
|
||||
@@ -224,7 +207,7 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head
|
||||
# Baseline (no SP)
|
||||
torch.multiprocessing.spawn(
|
||||
_run_attention_case,
|
||||
args=(1, base_port, input_file, baseline_file, num_heads, head_size, 1, "strict"),
|
||||
args=(1, base_init_method, input_file, baseline_file, num_heads, head_size, 1, "strict"),
|
||||
nprocs=1,
|
||||
)
|
||||
|
||||
@@ -233,7 +216,7 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head
|
||||
_run_attention_case,
|
||||
args=(
|
||||
sp_world_size,
|
||||
sp_port,
|
||||
sp_init_method,
|
||||
input_file,
|
||||
sp_file,
|
||||
num_heads,
|
||||
@@ -278,10 +261,8 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
|
||||
# rank0/1 -> 3+2=5, rank2/3 -> 3+2=5
|
||||
split_sizes = [3, 2, 3, 2]
|
||||
|
||||
base_port = _find_free_port()
|
||||
sp_port = _find_free_port()
|
||||
while sp_port == base_port:
|
||||
sp_port = _find_free_port()
|
||||
base_init_method = get_distributed_init_method()
|
||||
sp_init_method = get_distributed_init_method()
|
||||
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".npz") as f_in:
|
||||
input_file = f_in.name
|
||||
@@ -300,7 +281,19 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
|
||||
# Baseline (no SP)
|
||||
torch.multiprocessing.spawn(
|
||||
_run_attention_case,
|
||||
args=(1, base_port, input_file, baseline_file, num_heads, head_size, 1, "strict", 1, None, "mem_efficient"),
|
||||
args=(
|
||||
1,
|
||||
base_init_method,
|
||||
input_file,
|
||||
baseline_file,
|
||||
num_heads,
|
||||
head_size,
|
||||
1,
|
||||
"strict",
|
||||
1,
|
||||
None,
|
||||
"mem_efficient",
|
||||
),
|
||||
nprocs=1,
|
||||
)
|
||||
|
||||
@@ -309,7 +302,7 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None:
|
||||
_run_attention_case,
|
||||
args=(
|
||||
sp_world_size,
|
||||
sp_port,
|
||||
sp_init_method,
|
||||
input_file,
|
||||
sp_file,
|
||||
num_heads,
|
||||
|
||||
@@ -260,7 +260,7 @@ async def compare_single_vs_parallel(
|
||||
) -> None:
|
||||
"""Run the same prompts sequentially then in parallel and print a comparison."""
|
||||
|
||||
omni = AsyncOmni(model=model, diffusion_batch_size=batch_size)
|
||||
omni = AsyncOmni(model=model, max_num_seqs=batch_size, diffusion_batch_size=batch_size)
|
||||
try:
|
||||
await warmup(omni, WARMUP_PROMPTS)
|
||||
single_time = await run_single(omni, prompts)
|
||||
@@ -288,7 +288,7 @@ async def main(model: str, num_prompts: int, mode: str, batch_size: int = 1) ->
|
||||
await compare_single_vs_parallel(model, prompts, batch_size=batch_size)
|
||||
return
|
||||
|
||||
omni = AsyncOmni(model=model, diffusion_batch_size=batch_size)
|
||||
omni = AsyncOmni(model=model, max_num_seqs=batch_size, diffusion_batch_size=batch_size)
|
||||
try:
|
||||
await warmup(omni, WARMUP_PROMPTS)
|
||||
|
||||
@@ -524,7 +524,12 @@ if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="E2E diffusion concurrent benchmark / validation")
|
||||
parser.add_argument("--model", type=str, required=True, help="Model name or path")
|
||||
parser.add_argument("--num-prompts", type=int, default=8, help="Number of prompts to run")
|
||||
parser.add_argument("--batch-size", type=int, default=1, help="Diffusion batch size (1 = no batching)")
|
||||
parser.add_argument(
|
||||
"--batch-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Scheduler in-flight width (--max-num-seqs) and client diffusion_batch_size",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
choices=["batch", "single", "compare", "validate"],
|
||||
|
||||
+3
-6
@@ -365,13 +365,10 @@ class TestCacheDiTBackend:
|
||||
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter")
|
||||
@patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit")
|
||||
def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter):
|
||||
"""Test DreamID uses pipeline.transformer for cache enable/refresh.
|
||||
|
||||
NOTE: DreamID no longer has a custom enabler, so this tests against the generic path.
|
||||
"""
|
||||
def test_enable_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter):
|
||||
"""Generic cache path should pick up ``transformer.fused_blocks``."""
|
||||
mock_pipeline = Mock()
|
||||
mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline"
|
||||
mock_pipeline.__class__.__name__ = "FakeFusedBlocksPipeline"
|
||||
mock_pipeline.transformer = Mock()
|
||||
mock_pipeline.transformer.fused_blocks = Mock()
|
||||
mock_pipeline.transformer._cache_dit_adapter_config = CacheDiTAdapterConfig(
|
||||
|
||||
Loaded 100 of 1396 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user