From 47b18936766ea8ca97e030dce9c70f2e79315fa1 Mon Sep 17 00:00:00 2001 From: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com> Date: Wed, 19 Aug 2026 19:49:44 -0700 Subject: [PATCH] [Perf][CI] Reuse the Whisper judge worker across a test module (#6208) Signed-off-by: Shenglei Fu Co-authored-by: Hongsheng Liu --- tests/conftest.py | 1 + tests/helpers/fixtures/media.py | 24 +++ tests/helpers/media.py | 175 ++++++++++++---- tests/helpers/runtime.py | 21 +- tests/helpers/tests/test_media.py | 193 ++++++++++++++++-- .../tests/test_runtime_transcriber_release.py | 74 +++++++ 6 files changed, 428 insertions(+), 60 deletions(-) create mode 100644 tests/helpers/fixtures/media.py create mode 100644 tests/helpers/tests/test_runtime_transcriber_release.py diff --git a/tests/conftest.py b/tests/conftest.py index b453bbcbf..5c1276ec2 100644 --- a/tests/conftest.py +++ b/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", diff --git a/tests/helpers/fixtures/media.py b/tests/helpers/fixtures/media.py new file mode 100644 index 000000000..c34de093e --- /dev/null +++ b/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() diff --git a/tests/helpers/media.py b/tests/helpers/media.py index 8eb82fc1f..ed8ef035e 100644 --- a/tests/helpers/media.py +++ b/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", ] diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 4d06eb86a..56f7aaccc 100644 --- a/tests/helpers/runtime.py +++ b/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( diff --git a/tests/helpers/tests/test_media.py b/tests/helpers/tests/test_media.py index 4a8e5393d..8045f7a69 100644 --- a/tests/helpers/tests/test_media.py +++ b/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 diff --git a/tests/helpers/tests/test_runtime_transcriber_release.py b/tests/helpers/tests/test_runtime_transcriber_release.py new file mode 100644 index 000000000..6bc7d4221 --- /dev/null +++ b/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"