[Perf][CI] Use Whisper on GPU when memory permits (#5675)

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>
Signed-off-by: andyluo7 <andy.luo@amd.com>
Co-authored-by: andyluo7 <andy.luo@amd.com>

.buildkite/cuda/test-merge.yml

@@ -79,6 +79,7 @@ steps:
     steps:
       - label: "TTS · Qwen3-TTS CustomVoice Test"
         source_file_dependencies:
+          - tests/helpers/media.py
           - tests/e2e/online_serving/test_qwen3_tts_customvoice.py
           - tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
           - vllm_omni/model_executor/models/qwen3_tts/
@@ -101,6 +102,7 @@ steps:
 
       - label: "TTS · Qwen3-TTS Base Test"
         source_file_dependencies:
+          - tests/helpers/media.py
           - tests/e2e/online_serving/test_qwen3_tts_base.py
           - tests/e2e/offline_inference/test_qwen3_tts_base.py
           - vllm_omni/model_executor/models/qwen3_tts/
@@ -161,10 +163,11 @@ steps:
         commands:
           - pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5.py -m 'advanced_model and cuda' --run-level 'advanced_model'
         # h100_2, not h100_1: advanced_model asserts transcribe every audio
-        # response with Whisper, and tests/helpers/media.py only uses an
-        # accelerator for Whisper when a spare GPU exists. On a single card it
-        # grades on CPU: ~8-10 min per clip vs ~30 s, which alone overruns the
-        # 50-minute job timeout (build 2956).
+        # response with Whisper, and tests/helpers/media.py requires at least
+        # 16 GiB free VRAM on a device to run Whisper on GPU. When the model
+        # server leaves insufficient memory, Whisper falls back to CPU:
+        # ~8-10 min per clip vs ~30 s, which overruns the 50-minute job timeout
+        # (build 2956). h100_2 guarantees a second GPU with ample free VRAM.
         mirror_hardwares: h100_2
 
       - label: "Omni · MiniCPM-o 4.5 Duplex Test"

.buildkite/cuda/test-ready.yml

@@ -250,6 +250,7 @@ steps:
 
       - label: "TTS · Qwen3-TTS CustomVoice Test"
         source_file_dependencies:
+          - tests/helpers/media.py
           - tests/e2e/online_serving/test_qwen3_tts_customvoice.py
           - vllm_omni/model_executor/models/qwen3_tts/
           - vllm_omni/model_executor/models/common/qwen3_code_predictor.py

tests/helpers/media.py

@@ -31,6 +31,8 @@ from PIL import Image
 
 logger = logging.getLogger(__name__)
 
+_MIN_FREE_VRAM = 16 * 1024**3
+
 
 _synthetic_media_fallback_dir: Path | None = None
 
@@ -711,12 +713,16 @@ def _select_whisper_device() -> str:
 
     if current_omni_platform.is_available():
         n = current_omni_platform.get_device_count()
-        # Single-GPU runners (e.g. the L4 nightly): the model server already
-        # occupies device 0, and a Whisper model resident there competes with
-        # it for VRAM and OOMs. Only borrow an accelerator when a spare device
-        # exists; otherwise validate on CPU.
-        if n > 1:
-            device_index = n - 1
+        # Prefer the highest-index device with enough room for Whisper. The
+        # 16 GiB floor protects low-memory single-GPU runners such as the L4
+        # nightly, where the model server already occupies device 0, from the
+        # server-plus-Whisper OOM fixed in #3822. Keep the CPU fallback when no
+        # device has enough room.
+        for candidate_index in range(n - 1, -1, -1):
+            candidate_device = current_omni_platform.get_torch_device(candidate_index)
+            if current_omni_platform.get_free_memory(candidate_device) >= _MIN_FREE_VRAM:
+                device_index = candidate_index
+                break
 
     if device_index is None:
         return "cpu"

tests/helpers/tests/test_media.py

@@ -1,9 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import sys
 from concurrent.futures.process import BrokenProcessPool
 from contextlib import nullcontext
+from multiprocessing.context import BaseContext
 from types import SimpleNamespace
 
 import numpy as np
@@ -69,6 +70,7 @@ class _FakeExecutor:
 
     def __init__(self, outcomes=()):
         self._outcomes = list(outcomes)
+        self.init_kwargs: dict[str, object] = {}
         self.submitted: list[tuple] = []
         self.shutdown_calls = 0
 
@@ -113,6 +115,44 @@ def _reset_transcriber_singletons():
     media._WHISPER_MODELS.clear()
 
 
+@pytest.mark.parametrize(
+    ("free_memory_by_device", "expected_device", "expected_probe_order"),
+    [
+        ([media._MIN_FREE_VRAM], "cuda:0", [0]),
+        ([media._MIN_FREE_VRAM] * 3, "cuda:2", [2]),
+        ([8 * 1024**3, 20 * 1024**3, 24 * 1024**3], "cuda:2", [2]),
+        ([media._MIN_FREE_VRAM - 1], "cpu", [0]),
+        ([20 * 1024**3, 8 * 1024**3], "cuda:0", [1, 0]),
+        ([8 * 1024**3, 12 * 1024**3], "cpu", [1, 0]),
+    ],
+)
+def test_select_whisper_device_by_available_memory(
+    monkeypatch, free_memory_by_device, expected_device, expected_probe_order
+):
+    probed_devices = []
+
+    def get_free_memory(device):
+        index = int(device.rsplit(":", 1)[1])
+        probed_devices.append(index)
+        return free_memory_by_device[index]
+
+    platform = SimpleNamespace(
+        is_available=lambda: True,
+        get_device_count=lambda: len(free_memory_by_device),
+        get_torch_device=lambda index: f"cuda:{index}",
+        get_free_memory=get_free_memory,
+        set_device=lambda device: None,
+    )
+    monkeypatch.setitem(
+        sys.modules,
+        "vllm_omni.platforms",
+        SimpleNamespace(current_omni_platform=platform),
+    )
+
+    assert media._select_whisper_device() == expected_device
+    assert probed_devices == expected_probe_order
+
+
 def test_bytes_entrypoint_forwards_language_to_subprocess(monkeypatch, tmp_path):
     """Cover the two hops the tests above skip: bytes -> file -> executor.submit.
 
@@ -147,7 +187,9 @@ def test_parent_reuses_one_spawned_worker_across_calls(monkeypatch):
     assert len(created) == 1
     assert len(created[0].submitted) == 2
     assert created[0].init_kwargs["max_workers"] == 1
-    assert created[0].init_kwargs["mp_context"].get_start_method() == "spawn"
+    mp_context = created[0].init_kwargs["mp_context"]
+    assert isinstance(mp_context, BaseContext)
+    assert mp_context.get_start_method() == "spawn"
 
 
 def test_release_forces_a_new_process_pool(monkeypatch):