From 3ab726c126d760f6052c815b9bdced135cd86154 Mon Sep 17 00:00:00 2001 From: andyluo7 <43718156+andyluo7@users.noreply.github.com> Date: Sat, 3 Oct 2026 23:19:22 -0700 Subject: [PATCH] test: bound CPU vocoder reference threads (#8347) Signed-off-by: andyluo7 --- tests/conftest.py | 1 + tests/helpers/fixtures/cpu_threads.py | 52 +++++++++++ tests/helpers/tests/test_cpu_threads.py | 89 +++++++++++++++++++ .../test_cosyvoice3_f0_weight_norm.py | 1 + .../test_cosyvoice3_hift_weight_norm.py | 1 + .../test_cosyvoice3_incremental_hift.py | 7 +- .../test_qwen3_tts_incremental_decode.py | 26 +++--- 7 files changed, 160 insertions(+), 17 deletions(-) create mode 100644 tests/helpers/fixtures/cpu_threads.py create mode 100644 tests/helpers/tests/test_cpu_threads.py diff --git a/tests/conftest.py b/tests/conftest.py index fa55a23da..883967b1d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -23,6 +23,7 @@ bootstrap_vllm_layer_custom_op_modules() pytest_plugins = ( "tests.helpers.fixtures.config", "tests.helpers.fixtures.clean", + "tests.helpers.fixtures.cpu_threads", "tests.helpers.fixtures.log", "tests.helpers.fixtures.media", "tests.helpers.fixtures.pytest_collection", diff --git a/tests/helpers/fixtures/cpu_threads.py b/tests/helpers/fixtures/cpu_threads.py new file mode 100644 index 000000000..f03afd141 --- /dev/null +++ b/tests/helpers/fixtures/cpu_threads.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import time +from collections.abc import Callable, Generator +from typing import Protocol + +import pytest + + +class _TorchThreadAPI(Protocol): + def get_num_threads(self) -> int: ... + + def set_num_threads(self, num_threads: int) -> None: ... + + +def _single_threaded_cpu( + request: pytest.FixtureRequest, + load_torch: Callable[[], _TorchThreadAPI], +) -> Generator[None, None, None]: + """Implement the marker guard separately so its state handling is testable.""" + if request.node.get_closest_marker("cpu") is None: + yield + return + + torch = load_torch() + previous_threads = torch.get_num_threads() + started = time.perf_counter() + try: + torch.set_num_threads(1) + print(f"{request.node.nodeid}: CPU reference threads {previous_threads} -> {torch.get_num_threads()}") + yield + finally: + torch.set_num_threads(previous_threads) + print( + f"{request.node.nodeid}: CPU reference elapsed {time.perf_counter() - started:.3f}s; " + f"restored threads {torch.get_num_threads()}" + ) + + +@pytest.fixture +def single_threaded_cpu(request: pytest.FixtureRequest) -> Generator[None, None, None]: + """Bound opted-in CPU tests while leaving GPU-marked cases unchanged.""" + + def load_torch() -> _TorchThreadAPI: + import torch + + return torch + + yield from _single_threaded_cpu(request, load_torch) diff --git a/tests/helpers/tests/test_cpu_threads.py b/tests/helpers/tests/test_cpu_threads.py new file mode 100644 index 000000000..d467ea740 --- /dev/null +++ b/tests/helpers/tests/test_cpu_threads.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from tests.helpers.fixtures.cpu_threads import _single_threaded_cpu + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _ThreadController: + def __init__(self, num_threads: int, *, fail_when_bounding: bool = False): + self.num_threads = num_threads + self.fail_when_bounding = fail_when_bounding + self.transitions: list[int] = [] + + def get_num_threads(self) -> int: + return self.num_threads + + def set_num_threads(self, num_threads: int) -> None: + self.num_threads = num_threads + self.transitions.append(num_threads) + if self.fail_when_bounding and num_threads == 1: + raise RuntimeError("could not set thread count") + + +def _request(*, cpu_marked: bool): + marker = object() if cpu_marked else None + node = SimpleNamespace( + nodeid="tests/helpers/tests/test_cpu_threads.py::probe", + get_closest_marker=lambda name: marker if name == "cpu" else None, + ) + return SimpleNamespace(node=node) + + +def _finish_fixture(fixture) -> None: + with pytest.raises(StopIteration): + next(fixture) + + +def test_registered_fixture_caps_cpu_marked_test(single_threaded_cpu): + import torch + + assert torch.get_num_threads() == 1 + + +def test_single_threaded_cpu_caps_and_restores(capsys): + torch = _ThreadController(8) + fixture = _single_threaded_cpu(_request(cpu_marked=True), lambda: torch) + + next(fixture) + assert torch.num_threads == 1 + _finish_fixture(fixture) + + assert torch.num_threads == 8 + assert torch.transitions == [1, 8] + output = capsys.readouterr().out + assert "CPU reference threads 8 -> 1" in output + assert "restored threads 8" in output + + +def test_single_threaded_cpu_is_noop_without_cpu_marker(): + loaded = False + + def load_torch(): + nonlocal loaded + loaded = True + return _ThreadController(8) + + fixture = _single_threaded_cpu(_request(cpu_marked=False), load_torch) + next(fixture) + _finish_fixture(fixture) + + assert not loaded + + +def test_single_threaded_cpu_restores_if_bounding_raises(): + torch = _ThreadController(8, fail_when_bounding=True) + fixture = _single_threaded_cpu(_request(cpu_marked=True), lambda: torch) + + with pytest.raises(RuntimeError, match="could not set thread count"): + next(fixture) + + assert torch.num_threads == 8 + assert torch.transitions == [1, 8] diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_f0_weight_norm.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_f0_weight_norm.py index 2e149a815..f47a17e8e 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_f0_weight_norm.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_f0_weight_norm.py @@ -10,6 +10,7 @@ from torch.nn.utils import parametrize from vllm_omni.model_executor.models.cosyvoice3.code2wav_core import hifigan pytestmark = [ + pytest.mark.usefixtures("single_threaded_cpu"), pytest.mark.core_model, pytest.mark.filterwarnings("ignore:.*weight_norm.*:FutureWarning"), ] diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_weight_norm.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_weight_norm.py index 713d437ca..ea9608aa7 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_weight_norm.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_weight_norm.py @@ -10,6 +10,7 @@ from torch.nn.utils import parametrize from vllm_omni.model_executor.models.cosyvoice3.code2wav_core import hifigan pytestmark = [ + pytest.mark.usefixtures("single_threaded_cpu"), pytest.mark.core_model, pytest.mark.cpu, # The legacy weight_norm API warns on construction; the "legacy" params diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py index 1dd0799ef..940b8a476 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_incremental_hift.py @@ -31,9 +31,8 @@ from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( _wrapped_slice, ) from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav -from vllm_omni.platforms import current_omni_platform -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.usefixtures("single_threaded_cpu")] CHUNK_LEN = 24 TOTAL_MEL = 96 @@ -341,10 +340,6 @@ def test_incremental_hift_matches_reference_for_voiced_f0(config): @pytest.mark.parametrize("config", CONFIGS) @pytest.mark.parametrize("first_len,step_len", [(10, 6), (10, 3), (20, 3), (10, 1)]) -@pytest.mark.skipif( - current_omni_platform.is_rocm(), - reason="exhaustive CPU reference matrix exceeds the AMD CI time budget", -) def test_incremental_hift_matches_reference_for_voiced_f0_small_chunks(config, first_len, step_len): """Same as test_incremental_hift_matches_reference_for_voiced_f0, but with small, sub-receptive-field chunk sizes (down to 1 mel frame) instead of the uniform diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_incremental_decode.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_incremental_decode.py index 58cb04884..6de4e5193 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_incremental_decode.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_incremental_decode.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 copy from collections import Counter from types import SimpleNamespace +from typing import Any import pytest import torch @@ -22,7 +23,7 @@ from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts Qwen3TTSTokenizerV2DecoderTransformerModel, ) -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.usefixtures("single_threaded_cpu")] # Shared test helpers @@ -341,7 +342,7 @@ def test_incremental_request_audio_matches_full_decode(mode): (1, decoder.config.num_quantizers, prefix_frames + suffix_frames), ) suffix_codes = codes[..., prefix_frames:] - caches = {"prefix_frames": prefix_frames} + caches: dict[str, Any] = {"prefix_frames": prefix_frames} with torch.no_grad(): first = decoder(codes[..., : prefix_frames + 1], caches=caches) @@ -383,7 +384,7 @@ def test_single_frame_xvec_prefix_keeps_all_next_chunk_conv_frames(): decoder = Qwen3TTSTokenizerV2Decoder(config).eval() decoder._incremental_chunk_frames = 25 codes = torch.randint(0, config.codebook_size, (1, config.num_quantizers, 26)) - caches = {"prefix_frames": 0} + caches: dict[str, Any] = {"prefix_frames": 0} with torch.no_grad(): decoder(codes[..., :1], caches=caches) @@ -421,7 +422,7 @@ def test_xvec_rolling_matches_truncated_suffix_window(): decoder = Qwen3TTSTokenizerV2Decoder(config).eval() decoder._incremental_chunk_frames = 25 codes = torch.randint(0, config.codebook_size, (1, config.num_quantizers, 101)) - caches = {"prefix_frames": 0} + caches: dict[str, Any] = {"prefix_frames": 0} with torch.no_grad(): decoder(codes[..., :1], caches=caches) @@ -470,7 +471,7 @@ def test_icl_rolling_matches_reference_plus_truncated_suffix_window(): ) prefix_codes = codes[..., :prefix_frames] suffix_codes = codes[..., prefix_frames:] - caches = {"prefix_frames": prefix_frames} + caches: dict[str, Any] = {"prefix_frames": prefix_frames} with torch.no_grad(): decoder(codes[..., : prefix_frames + 1], caches=caches) @@ -975,7 +976,7 @@ def test_batched_chunked_decode_groups_exact_phases(monkeypatch): prefix_frames = 48 codes_list = [torch.zeros(1, 2, prefix_frames + 1)] - caches = [{"prefix_frames": prefix_frames}] + caches: list[dict[str, Any]] = [{"prefix_frames": prefix_frames}] for previous, _suffix_frames, cached_frames in ( (1, 26, 1), (26, 51, 26), @@ -1064,7 +1065,7 @@ def test_missing_graph_phase_falls_back_instead_of_returning_empty(monkeypatch): wrapper._record_graph_fallback = lambda *_args: None codes = [torch.zeros(1, 2, 25)] - caches = [{}] + caches: list[dict[str, Any]] = [{}] assert wrapper._decode_icl_prefix_batch(codes, caches) is None assert wrapper._decode_suffix_batch("xvec", 50, codes, caches, [25]) is None @@ -1078,9 +1079,12 @@ def test_xvec_first_chunk_uses_prefixless_initializer(monkeypatch): wrapper._icl_previous_frames_by_target = {26: 1} wrapper._xvec_previous_frames_by_target = {26: 1} calls = [] - wrapper.decoder = _decoder_stub( - _decode_xvec_first_chunk=lambda codes, cache: calls.append("xvec") or codes[:, :1, :], - ) + + def decode_first_chunk(codes, cache): + calls.append("xvec") + return codes[:, :1, :] + + wrapper.decoder = _decoder_stub(_decode_xvec_first_chunk=decode_first_chunk) output = wrapper._batched_request_decode( [torch.zeros(1, 2, 25)],