[Perf][CI] Reuse the Whisper judge worker across a test module (#6208)

Signed-off-by: Shenglei Fu <sfu@confluent.io>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/conftest.py

@@ -21,6 +21,7 @@ pytest_plugins = (
     "tests.helpers.fixtures.config",
     "tests.helpers.fixtures.env",
     "tests.helpers.fixtures.log",
+    "tests.helpers.fixtures.media",
     "tests.helpers.fixtures.run_args",
     "tests.helpers.fixtures.runtime",
     "tests.helpers.fixtures.speaker_cache",

tests/helpers/fixtures/media.py

@@ -0,0 +1,24 @@
+"""Fixtures for the reused Whisper transcription worker."""
+
+import sys
+
+import pytest
+
+
+@pytest.fixture(autouse=True, scope="module")
+def release_audio_transcriber_after_module():
+    """Free the transcription worker's device memory once a test module is done.
+
+    A backstop. The server and runner fixtures already release the worker when
+    each of their instances tears down, which is what keeps Whisper off the
+    device while the next parametrized server initializes. This covers the
+    remaining case: a module that transcribes without using those fixtures.
+
+    Reached through ``sys.modules`` because ``tests/conftest.py`` deliberately
+    keeps ``tests.helpers.media`` lazy: a module that never transcribed should
+    not pay to import it.
+    """
+    yield
+    media = sys.modules.get("tests.helpers.media")
+    if media is not None:
+        media.release_audio_transcriber()

tests/helpers/media.py

@@ -3,7 +3,6 @@
 import atexit
 import base64
 import concurrent.futures
-import gc
 import hashlib
 import io
 import logging
@@ -14,7 +13,9 @@ import random
 import re
 import subprocess
 import tempfile
+import threading
 import time
+from concurrent.futures.process import BrokenProcessPool
 from contextlib import contextmanager
 from pathlib import Path
 from typing import Any
@@ -670,59 +671,152 @@ def _serialize_whisper_model_download(model_size: str = "small"):
         f.close()
 
 
-def _whisper_transcribe_in_current_process(
-    output_path: str, model_size: str = "small", language: str | None = None
-) -> str:
-    import whisper
-
+def _select_whisper_device() -> str:
     device_index = None
     from vllm_omni.platforms import current_omni_platform
 
     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. Loading Whisper there, once per concurrent
-        # request, competes for VRAM and OOMs. Only borrow an accelerator
-        # when a spare device exists; otherwise validate on CPU.
+        # 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
 
-    if device_index is not None:
-        torch_device = current_omni_platform.get_torch_device(device_index)
-        current_omni_platform.set_device(torch_device)
-        device = str(torch_device)
-        use_accelerator = True
-    else:
-        use_accelerator = False
-        device = "cpu"
+    if device_index is None:
+        return "cpu"
 
-    with _serialize_whisper_model_download(model_size):
-        model = whisper.load_model(model_size, device=device)
-    try:
-        text = model.transcribe(
-            output_path,
-            temperature=0.0,
-            word_timestamps=True,
-            condition_on_previous_text=False,
-            # None keeps whisper's auto-detection. Do not default this to a
-            # language: callers include non-English audio tests.
-            language=language,
-        )["text"]
-    finally:
-        del model
-        gc.collect()
-        if use_accelerator:
-            current_omni_platform.synchronize()
-            current_omni_platform.empty_cache()
+    torch_device = current_omni_platform.get_torch_device(device_index)
+    current_omni_platform.set_device(torch_device)
+    return str(torch_device)
+
+
+# Populated in the transcription worker, not in the pytest process.
+_WHISPER_MODELS: dict[str, Any] = {}
+
+
+def _get_whisper_model(model_size: str) -> Any:
+    model = _WHISPER_MODELS.get(model_size)
+    if model is None:
+        import whisper
+
+        # The device is picked on first load and the model stays on it for the
+        # worker's lifetime: the current server or runner fixture instance, or
+        # the test module for callers that transcribe without those fixtures.
+        device = _select_whisper_device()
+        with _serialize_whisper_model_download(model_size):
+            model = whisper.load_model(model_size, device=device)
+        _WHISPER_MODELS[model_size] = model
+    return model
+
+
+def _whisper_transcribe_in_current_process(
+    output_path: str, model_size: str = "small", language: str | None = None
+) -> str:
+    model = _get_whisper_model(model_size)
+    text = model.transcribe(
+        output_path,
+        temperature=0.0,
+        word_timestamps=True,
+        condition_on_previous_text=False,
+        # None keeps whisper's auto-detection. Do not default this to a
+        # language: callers include non-English audio tests.
+        language=language,
+    )["text"]
     return text or ""
 
 
+# Serializes a whole submit->result->cleanup on the parent side, so at most one
+# call is ever in flight on the single-worker pool. The child already runs
+# transcriptions serially, so this costs no throughput; it lets a failed call
+# discard the worker without disrupting another caller's in-flight future.
+_TRANSCRIBER_CALL_LOCK = threading.Lock()
+# Guards the _TRANSCRIBER pointer itself.
+_TRANSCRIBER_LOCK = threading.Lock()
+_TRANSCRIBER: concurrent.futures.ProcessPoolExecutor | None = None
+
+
+def _get_transcriber() -> concurrent.futures.ProcessPoolExecutor:
+    global _TRANSCRIBER
+    with _TRANSCRIBER_LOCK:
+        if _TRANSCRIBER is None:
+            ctx = multiprocessing.get_context("spawn")
+            _TRANSCRIBER = concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx)
+        return _TRANSCRIBER
+
+
+def _discard_transcriber(executor: concurrent.futures.ProcessPoolExecutor) -> None:
+    """Drop ``executor``, but only while it is still the current one.
+
+    Identity-checked so a stale reference can never shut down a newer worker that
+    was installed after ``executor`` was replaced.
+    """
+    global _TRANSCRIBER
+    with _TRANSCRIBER_LOCK:
+        if _TRANSCRIBER is not executor:
+            return
+        _TRANSCRIBER = None
+    # Joining the worker can block; do it outside the lock.
+    executor.shutdown(wait=True)
+
+
+def release_audio_transcriber() -> None:
+    """Shut the transcription worker down, freeing the device memory its models hold.
+
+    Called when a server or runner fixture instance tears down, so that the next
+    one -- including the next parametrization inside the same test module -- does
+    not initialize its model while Whisper still occupies the device. A module
+    teardown fixture repeats it for tests that transcribe without those fixtures.
+
+    Takes the call lock, so it waits for any in-flight transcription to finish
+    rather than shutting the worker down underneath it.
+    """
+    global _TRANSCRIBER
+    with _TRANSCRIBER_CALL_LOCK:
+        with _TRANSCRIBER_LOCK:
+            executor, _TRANSCRIBER = _TRANSCRIBER, None
+        if executor is not None:
+            executor.shutdown(wait=True)
+
+
 def convert_audio_file_to_text(output_path: str, model_size: str = "small", language: str | None = None) -> str:
-    """Convert an audio file to text in an isolated subprocess."""
-    ctx = multiprocessing.get_context("spawn")
-    with concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx) as executor:
-        future = executor.submit(_whisper_transcribe_in_current_process, output_path, model_size, language)
-        return future.result()
+    """Convert an audio file to text in a reused, isolated subprocess.
+
+    The worker outlives the call so its Whisper model is loaded once rather than
+    once per transcription. The call lock serializes callers onto the single
+    worker (the child already transcribes serially, so this costs no throughput),
+    which lets a failed call tear the worker down without racing another caller.
+    The worker caches one model per size it is asked for, so a run that escalates
+    to a stronger ASR keeps both models resident until release.
+
+    Any failure discards the worker, restoring the failure isolation of the old
+    one-process-per-call design: a failure (a ``torch`` OOM, or a
+    ``KeyboardInterrupt``/``SystemExit`` transported through the future or
+    raised while ``result()`` blocks) propagates unchanged after the worker --
+    and its resident model -- is torn down, and a dead worker
+    (``BrokenProcessPool``) is additionally retried once.
+    """
+    with _TRANSCRIBER_CALL_LOCK:
+        for attempt in range(2):
+            executor = _get_transcriber()
+            try:
+                return executor.submit(
+                    _whisper_transcribe_in_current_process, output_path, model_size, language
+                ).result()
+            except BrokenProcessPool:
+                _discard_transcriber(executor)
+                if attempt == 1:
+                    raise
+            except BaseException:
+                # Any other failure -- a task exception (a torch OOM included), or
+                # a KeyboardInterrupt/SystemExit transported through the future or
+                # interrupting result() -- leaves the worker and its model
+                # resident; drop it so it cannot contaminate later calls, matching
+                # the old per-call teardown. Do not retry.
+                _discard_transcriber(executor)
+                raise
+    raise AssertionError("unreachable")
 
 
 def convert_audio_bytes_to_text(raw_bytes: bytes, model_size: str = "small", language: str | None = None) -> str:
@@ -749,4 +843,5 @@ __all__ = [
     "get_asset_path",
     "load_test_audio_data_url",
     "preprocess_text",
+    "release_audio_transcriber",
 ]

tests/helpers/runtime.py

@@ -16,6 +16,7 @@ import tempfile
 import threading
 import time
 from collections.abc import Generator
+from contextlib import contextmanager
 from dataclasses import asdict, dataclass
 from io import BytesIO
 from pathlib import Path
@@ -50,6 +51,7 @@ from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
 from tests.helpers.media import (
     _merge_base64_audio_to_segment,
     decode_b64_image,
+    release_audio_transcriber,
 )
 from tests.model_tests.diffusion.utils import resolve_tiny_model_path
 from vllm_omni.config.stage_config import resolve_deploy_yaml
@@ -3143,6 +3145,21 @@ class OmniRunnerHandler:
 # ---------------------------------------------------------------------------
 
 
+@contextmanager
+def _whisper_device_free_around():
+    """Keep the Whisper judge off the accelerator around a server/runner lifecycle.
+
+    Released on entry so the next instance -- e.g. the next parametrization in
+    the same module -- initializes on a clean device, and again on exit
+    (including when a test raises) so it does not linger for the instance after.
+    """
+    release_audio_transcriber()
+    try:
+        yield
+    finally:
+        release_audio_transcriber()
+
+
 def iter_omni_server(
     request: Any,
     run_level: str,
@@ -3152,7 +3169,7 @@ def iter_omni_server(
     """Start/stop an Omni HTTP server; used by ``omni_server`` / ``omni_server_function`` fixtures."""
     from tests.helpers.stage_config import stage_config_path_for_run_level
 
-    with omni_fixture_lock:
+    with omni_fixture_lock, _whisper_device_free_around():
         params: OmniServerParams = request.param
         # For now, when a tiny model is substituted, we preserve the original model
         # name via --served-model-name (so that the server still accepts requests with
@@ -3240,7 +3257,7 @@ def iter_omni_runner(
     """Yield an :class:`OmniRunner`; used by ``omni_runner`` / ``omni_runner_function`` fixtures."""
     from tests.helpers.stage_config import stage_config_path_for_run_level
 
-    with omni_fixture_lock:
+    with omni_fixture_lock, _whisper_device_free_around():
         param = request.param
         if not isinstance(param, (tuple, list)) or len(param) not in (2, 3):
             raise ValueError(

tests/helpers/tests/test_media.py

@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import sys
+from concurrent.futures.process import BrokenProcessPool
+from contextlib import nullcontext
 from types import SimpleNamespace
 
 import numpy as np
@@ -22,7 +24,11 @@ def _stub_transcribe(monkeypatch) -> dict:
             captured.update(kwargs)
             return {"text": "London"}
 
-    monkeypatch.setitem(sys.modules, "whisper", SimpleNamespace(load_model=lambda size, device=None: FakeModel()))
+    monkeypatch.setitem(
+        sys.modules,
+        "whisper",
+        SimpleNamespace(load_model=lambda size, device=None: FakeModel()),
+    )
     # Keep the unit test off the accelerator probe so it stays hermetic.
     monkeypatch.setitem(
         sys.modules,
@@ -50,32 +56,183 @@ def test_transcribe_defaults_to_auto_language(monkeypatch):
     assert captured.get("language") is None
 
 
+class _FakeExecutor:
+    """Stand-in for ProcessPoolExecutor recording what was submitted and shut down."""
+
+    def __init__(self, outcomes=()):
+        self._outcomes = list(outcomes)
+        self.submitted: list[tuple] = []
+        self.shutdown_calls = 0
+
+    def submit(self, fn, *args):
+        self.submitted.append((fn, args))
+        outcome = self._outcomes.pop(0) if self._outcomes else "London"
+
+        def _result():
+            if isinstance(outcome, BaseException):
+                raise outcome
+            return outcome
+
+        return SimpleNamespace(result=_result)
+
+    def shutdown(self, wait=True):
+        self.shutdown_calls += 1
+
+
+def _patch_executors(monkeypatch, outcomes_per_executor=()) -> list[_FakeExecutor]:
+    """Replace the pool constructor and hand back the fakes it produces, in creation order."""
+    created: list[_FakeExecutor] = []
+
+    def factory(*args, **kwargs):
+        index = len(created)
+        outcomes = outcomes_per_executor[index] if index < len(outcomes_per_executor) else ()
+        executor = _FakeExecutor(outcomes)
+        executor.init_kwargs = kwargs
+        created.append(executor)
+        return executor
+
+    monkeypatch.setattr(media.concurrent.futures, "ProcessPoolExecutor", factory)
+    return created
+
+
+@pytest.fixture(autouse=True)
+def _reset_transcriber_singletons():
+    """The worker and its model cache are module-level state; do not leak them across tests."""
+    media.release_audio_transcriber()
+    media._WHISPER_MODELS.clear()
+    yield
+    media.release_audio_transcriber()
+    media._WHISPER_MODELS.clear()
+
+
 def test_bytes_entrypoint_forwards_language_to_subprocess(monkeypatch, tmp_path):
     """Cover the two hops the tests above skip: bytes -> file -> executor.submit.
 
     The real call crosses a spawn ProcessPoolExecutor, so capture what gets
     submitted rather than what the worker eventually does.
     """
-    submitted: dict = {}
-
-    class FakeExecutor:
-        def __enter__(self):
-            return self
-
-        def __exit__(self, *exc):
-            return False
-
-        def submit(self, fn, *args):
-            submitted["fn"] = fn
-            submitted["args"] = args
-            return SimpleNamespace(result=lambda: "London")
-
-    monkeypatch.setattr(media.concurrent.futures, "ProcessPoolExecutor", lambda *a, **kw: FakeExecutor())
+    created = _patch_executors(monkeypatch)
 
     wav = tmp_path / "clip.wav"
     sf.write(wav, np.zeros(2400, dtype=np.float32), 24000)
 
     assert media.convert_audio_bytes_to_text(wav.read_bytes(), "small", "en") == "London"
 
-    assert submitted["fn"] is media._whisper_transcribe_in_current_process
-    assert submitted["args"][1:] == ("small", "en")
+    fn, args = created[0].submitted[0]
+    assert fn is media._whisper_transcribe_in_current_process
+    assert args[1:] == ("small", "en")
+
+
+def test_parent_reuses_one_spawned_worker_across_calls(monkeypatch):
+    """Two calls share one spawned worker.
+
+    Pins the parent-side contract: the pool is constructed once, with
+    ``spawn`` (CUDA is not fork-safe) and ``max_workers=1`` (one resident
+    model). Reuse of the worker's *loaded model* is pinned by the worker-side
+    cache test below, which is where the startup cost actually lives.
+    """
+    created = _patch_executors(monkeypatch)
+
+    media.convert_audio_file_to_text("/tmp/a.wav")
+    media.convert_audio_file_to_text("/tmp/b.wav")
+
+    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"
+
+
+def test_release_forces_a_new_process_pool(monkeypatch):
+    created = _patch_executors(monkeypatch)
+
+    media.convert_audio_file_to_text("/tmp/a.wav")
+    media.release_audio_transcriber()
+    media.convert_audio_file_to_text("/tmp/b.wav")
+
+    assert len(created) == 2
+    assert created[0].shutdown_calls == 1
+
+
+def test_worker_loads_each_model_size_once(monkeypatch):
+    """The worker-side cache is what removes the repeated whisper.load_model cost."""
+    loaded: list[str] = []
+
+    class FakeModel:
+        def transcribe(self, path, **kwargs):
+            return {"text": "London"}
+
+    def load_model(size, device=None):
+        loaded.append(size)
+        return FakeModel()
+
+    monkeypatch.setitem(sys.modules, "whisper", SimpleNamespace(load_model=load_model))
+    monkeypatch.setitem(
+        sys.modules,
+        "vllm_omni.platforms",
+        SimpleNamespace(current_omni_platform=SimpleNamespace(is_available=lambda: False)),
+    )
+    monkeypatch.setattr(media, "_serialize_whisper_model_download", lambda model_size: nullcontext())
+
+    media._whisper_transcribe_in_current_process("/tmp/a.wav", "small")
+    media._whisper_transcribe_in_current_process("/tmp/b.wav", "small")
+
+    assert loaded == ["small"]
+
+    # The ASR escalation path asks the same worker for a stronger model.
+    media._whisper_transcribe_in_current_process("/tmp/c.wav", "large-v3")
+
+    assert loaded == ["small", "large-v3"]
+
+
+@pytest.mark.parametrize(
+    ("second_outcome", "recovers"),
+    [(["London"], True), ([BrokenProcessPool("second")], False)],
+)
+def test_broken_pool_is_discarded_and_retried_once(monkeypatch, second_outcome, recovers):
+    """A dead worker is discarded and the call retried exactly once.
+
+    The retry either succeeds, or -- if the replacement also dies -- fails with
+    both pools discarded and no executor left installed.
+    """
+    created = _patch_executors(
+        monkeypatch,
+        outcomes_per_executor=([BrokenProcessPool("first")], second_outcome),
+    )
+
+    if recovers:
+        assert media.convert_audio_file_to_text("/tmp/a.wav") == "London"
+    else:
+        with pytest.raises(BrokenProcessPool):
+            media.convert_audio_file_to_text("/tmp/a.wav")
+        assert created[1].shutdown_calls == 1
+        assert media._TRANSCRIBER is None
+
+    assert len(created) == 2  # discarded the first, retried on a fresh one
+    assert created[0].shutdown_calls == 1
+
+
+@pytest.mark.parametrize(
+    "error",
+    [RuntimeError("CUDA out of memory"), KeyboardInterrupt()],
+    ids=["Exception", "BaseException"],
+)
+def test_transcription_failure_propagates_and_discards_the_worker(monkeypatch, error):
+    """Any failure -- an Exception or a BaseException like KeyboardInterrupt --
+    propagates AND tears the worker down.
+
+    It arrives as the task's own exception, not ``BrokenProcessPool``, but the
+    worker -- and its resident model -- is still discarded, restoring the old
+    one-process-per-call isolation. It is not retried; the next call builds a
+    fresh worker.
+    """
+    created = _patch_executors(monkeypatch, outcomes_per_executor=([error],))
+
+    with pytest.raises(type(error)):
+        media.convert_audio_file_to_text("/tmp/a.wav")
+
+    assert len(created) == 1  # not retried
+    assert created[0].shutdown_calls == 1
+    assert media._TRANSCRIBER is None
+
+    assert media.convert_audio_file_to_text("/tmp/b.wav") == "London"
+    assert len(created) == 2  # a fresh worker for the next call

tests/helpers/tests/test_runtime_transcriber_release.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""The server/runner fixtures free the Whisper judge around each instance.
+
+iter_omni_server and iter_omni_runner both wrap their body in the same
+_whisper_device_free_around() context manager, so one representative path
+exercises its enter/exit and exception behavior.
+"""
+
+import threading
+from types import SimpleNamespace
+
+import pytest
+
+from tests.helpers import runtime
+from tests.helpers.runtime import OmniServerParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def events(monkeypatch) -> list[str]:
+    """Record construction and release in order."""
+    calls: list[str] = []
+
+    class FakeServer:
+        def __init__(self, *args, **kwargs):
+            calls.append("construct")
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *exc):
+            return False
+
+    monkeypatch.setattr(runtime, "OmniServer", FakeServer)
+    monkeypatch.setattr(runtime, "release_audio_transcriber", lambda: calls.append("release"))
+    monkeypatch.setattr(
+        "tests.helpers.stage_config.stage_config_path_for_run_level",
+        lambda path, run_level: None,
+    )
+    return calls
+
+
+def _server_generator():
+    request = SimpleNamespace(
+        param=OmniServerParams(model="fake-model"),
+        node=SimpleNamespace(get_closest_marker=lambda name: None),
+    )
+    return runtime.iter_omni_server(request, "core_model", "", threading.Lock())
+
+
+def test_releases_before_construction_and_after_teardown(events):
+    generator = _server_generator()
+    generator.send(None)
+
+    assert events == ["release", "construct"], "the device must be clear before the model loads"
+
+    with pytest.raises(StopIteration):
+        generator.send(None)
+
+    assert events == ["release", "construct", "release"]
+
+
+def test_failing_test_still_releases(events):
+    """A test that raises must not strand Whisper on the device."""
+    generator = _server_generator()
+    generator.send(None)
+
+    with pytest.raises(RuntimeError, match="assertion blew up"):
+        generator.throw(RuntimeError("assertion blew up"))
+
+    assert events[-1] == "release"