test: bound CPU vocoder reference threads (#8347)

Signed-off-by: andyluo7 <andy.luo@amd.com>

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",

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)

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]

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"),
 ]

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

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

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)],