diff --git a/funasr/auto/auto_model.py b/funasr/auto/auto_model.py index 16eacfbe45..82d4c226ee 100644 --- a/funasr/auto/auto_model.py +++ b/funasr/auto/auto_model.py @@ -53,6 +53,110 @@ def _resolve_ncpu(config, fallback=4): return max(value, 1) +# The active BLAS limit, if any. Held for the life of the process rather than +# scoped to a call -- see `_limit_blas_threads` for why. +_blas_thread_limiter = None +_blas_thread_limit = None + + +def _current_blas_threads(): + """Thread count currently configured for BLAS, or None if unreadable.""" + try: + import threadpoolctl + except ImportError: + return None + counts = { + pool["num_threads"] + for pool in threadpoolctl.threadpool_info() + if pool["user_api"] == "blas" + } + return counts.pop() if len(counts) == 1 else None + + +def _limit_blas_threads(ncpu): + """Apply `ncpu` to BLAS as well as to torch. + + `ncpu` documents itself as the thread count for "CPU 内部操作并行性", but + it only ever reached `torch.set_num_threads`. The speaker-clustering path + does not run on torch at all: `scipy.linalg.eigh` on the affinity + Laplacian goes to BLAS, whose own default is one thread per core. So on a + many-core host the CPU is saturated by a library `ncpu` never touched. + Measured on a 64-core host with a 357-row Laplacian: + + BLAS threads wall CPU + 64 1.20s 68.9 core-seconds + 4 0.05s 2.8 core-seconds + 1 0.03s 1.8 core-seconds + + At a few hundred rows the parallel driver spends its time synchronising + rather than computing, so bounding it is a latency win as well as a CPU + one. + + The update follows `torch.set_num_threads`: the newest `ncpu` wins, and no + set is issued when BLAS already sits at that value. Both settings are + process-wide, so constructing a second `AutoModel` applies its `ncpu` to + any model already built in this process -- including ones without speaker + clustering. That is the same contract `torch.set_num_threads` already + has; the difference is only that `ncpu` used to leave BLAS untouched. + + The limit is not released once applied. It is process-global, so a scoped + enter/exit pair per request could be interleaved by overlapping requests + and leave the setting unbalanced, and FunASR reaches BLAS only through + small one-shot operations. Doing nothing when threadpoolctl is missing + leaves BLAS at its own default rather than failing model construction. + """ + global _blas_thread_limiter, _blas_thread_limit + + current = _current_blas_threads() + if current is not None and current == ncpu: + return + + # `threadpool_limits` records the OpenMP pools as well as the BLAS ones + # even when entered with `user_api="blas"`, and its `__exit__` restores + # every pool it captured. Dropping the previous limiter would therefore + # roll torch's thread count back to whatever it was when that limiter was + # created -- undoing the `torch.set_num_threads(ncpu)` that `build_model` + # has already performed. Save and restore it across the swap so replacing + # the BLAS limit cannot touch an OpenMP setting it does not own. + torch_threads = torch.get_num_threads() + + if _blas_thread_limiter is not None: + _blas_thread_limiter.__exit__(None, None, None) + _blas_thread_limiter = None + + try: + import threadpoolctl + except ImportError: + logging.info( + "threadpoolctl is not installed, so `ncpu` cannot be applied to " + "BLAS; speaker clustering may use one thread per core." + ) + return + + _blas_thread_limiter = threadpoolctl.threadpool_limits( + limits=ncpu, user_api="blas" + ) + _blas_thread_limiter.__enter__() + _blas_thread_limit = ncpu + + if torch.get_num_threads() != torch_threads: + torch.set_num_threads(torch_threads) + + # BLAS pools are only visible to threadpoolctl once a library backed by + # them has been loaded. funasr imports such a library at package import, + # so they are present by the time a model is built -- but confirm rather + # than assume, because a dropped limit here would silently restore the + # CPU saturation this is meant to prevent. + applied = _current_blas_threads() + if applied is not None and applied != ncpu: + logging.warning( + "ncpu=%s could not be applied to BLAS (it reports %s); " + "speaker clustering may use more threads than requested.", + ncpu, + applied, + ) + + def _join_vad_texts(texts): """Remove rich tags and join VAD text without adding spaces between Chinese chunks.""" cleaned = [re.sub(r"<\|[^|]*\|>", "", text).strip() for text in texts] @@ -570,6 +674,7 @@ def build_model(**kwargs): kwargs["ncpu"] = ncpu if torch.get_num_threads() != ncpu: torch.set_num_threads(ncpu) + _limit_blas_threads(ncpu) # build tokenizer tokenizer = kwargs.get("tokenizer", None) diff --git a/funasr/models/campplus/cluster_backend.py b/funasr/models/campplus/cluster_backend.py index a369d626d4..759376f14d 100644 --- a/funasr/models/campplus/cluster_backend.py +++ b/funasr/models/campplus/cluster_backend.py @@ -104,12 +104,23 @@ def get_laplacian(self, M): def get_spec_embs(self, L, k_oracle=None): """Get spec embs. - + Args: L: TODO. k_oracle: TODO. """ - lambdas, eig_vecs = scipy.linalg.eigh(L) + # Only the leading eigenpairs are ever used: the speaker count comes + # from the gaps between the first ``max_num_spks + 1`` eigenvalues, and + # the embedding keeps the first ``num_of_spk`` eigenvectors. Producing + # the full spectrum costs O(N^3) and throws all but a handful away. + n_eig = self.max_num_spks + 1 + if k_oracle is not None: + n_eig = max(n_eig, int(k_oracle)) + n_eig = min(n_eig, L.shape[0]) + + lambdas, eig_vecs = scipy.linalg.eigh( + L, subset_by_index=[0, n_eig - 1] + ) if k_oracle is not None: num_of_spk = k_oracle diff --git a/setup.py b/setup.py index 762784d1b6..561da526d2 100644 --- a/setup.py +++ b/setup.py @@ -12,6 +12,7 @@ "install": [ # Core "scipy>=1.4.1", + "threadpoolctl>=3.0.0", "librosa", "soundfile>=0.12.1", # numpy 2.x is supported; deprecated aliases (np.float, np.int) were diff --git a/tests/test_auto_model_ncpu.py b/tests/test_auto_model_ncpu.py new file mode 100644 index 0000000000..d2982951b2 --- /dev/null +++ b/tests/test_auto_model_ncpu.py @@ -0,0 +1,218 @@ +"""Regression cover for applying `ncpu` to BLAS as well as to torch. + +Both settings are process-wide and follow `torch.set_num_threads` semantics. +The helper is driven directly for the update rule, and through the real +`AutoModel.build_model` path for the interaction between them -- swapping the +BLAS limiter touches torch's OpenMP pool as a side effect, which only shows up +across a sequence of builds. +""" + +import numpy as np +import pytest +import scipy.linalg +import torch +import threadpoolctl + +from funasr.auto.auto_model import ( + _current_blas_threads, + _limit_blas_threads, +) + + +@pytest.fixture(autouse=True) +def _restore_thread_settings(): + """Leave both thread settings where the test found them. + + The BLAS limit is deliberately process-wide and never released, so without + this a test would leak its thread count into every later test -- and the + torch count needs restoring for the same reason, since tests here set it + to make their assertions readable. + """ + original_blas = _current_blas_threads() + original_torch = torch.get_num_threads() + yield + if original_blas is not None: + _limit_blas_threads(original_blas) + torch.set_num_threads(original_torch) + + +def _require_blas_pools(): + if _current_blas_threads() is None: + pytest.skip("no BLAS pool is loaded in this environment") + + +class _ProbeModel(torch.nn.Module): + """Smallest thing `build_model` accepts, so no weights are downloaded.""" + + def __init__(self, **kwargs): + super().__init__() + self.linear = torch.nn.Linear(2, 2) + + +@pytest.fixture +def probe_model(): + from funasr.register import tables + + name = "ncpu-blas-probe" + added = name not in tables.model_classes + if added: + tables.model_classes[name] = _ProbeModel + yield name + if added: + tables.model_classes.pop(name, None) + + +def _build(probe_model, ncpu): + from funasr.auto.auto_model import AutoModel + + AutoModel( + model=probe_model, + model_conf={}, + device="cpu", + ncpu=ncpu, + disable_update=True, + ) + return torch.get_num_threads(), _current_blas_threads() + + +def test_build_model_sets_both_settings_in_both_orders(probe_model): + """End-to-end through `build_model`, where the settings are applied. + + Replacing the BLAS limiter drops the previous `threadpool_limits`, and that + object restores the OpenMP pools as well -- so without care the swap rolls + torch's thread count back to its value at the previous build, undoing the + `torch.set_num_threads(ncpu)` on the line above. Asserting only on the + helper misses this; a sequence of builds is what surfaces it. + """ + _require_blas_pools() + + for order in ((8, 1), (1, 8), (4, 2, 7)): + for ncpu in order: + torch_threads, blas_threads = _build(probe_model, ncpu) + assert blas_threads == ncpu, ( + "order %s: BLAS settled on %s at ncpu=%s" + % (order, blas_threads, ncpu) + ) + assert torch_threads == ncpu, ( + "order %s: torch settled on %s at ncpu=%s" + % (order, torch_threads, ncpu) + ) + + +def test_build_model_without_speaker_clustering(probe_model): + """`ncpu` reaches BLAS whether or not clustering is configured.""" + _require_blas_pools() + + torch_threads, blas_threads = _build(probe_model, 3) + assert blas_threads == 3 + assert torch_threads == 3 + + +def test_latest_ncpu_wins(): + """Later construction must be able to lower the limit. + + The first revision returned early once a limit existed, so ncpu=1 after + ncpu=8 left BLAS at 8 -- exactly the case this change exists to fix. + """ + _require_blas_pools() + + _limit_blas_threads(8) + assert _current_blas_threads() == 8 + + _limit_blas_threads(1) + assert _current_blas_threads() == 1, ( + "a later, lower ncpu must take effect" + ) + + +def test_latest_ncpu_wins_in_both_orders(): + _require_blas_pools() + + for first, second in ((8, 1), (1, 8), (4, 2), (2, 4)): + _limit_blas_threads(first) + assert _current_blas_threads() == first + _limit_blas_threads(second) + assert _current_blas_threads() == second, ( + "order %s -> %s did not settle on the newest value" + % (first, second) + ) + + +def test_repeated_same_value_is_stable(): + """Re-setting the current value must be a no-op, not a nested limit.""" + _require_blas_pools() + + _limit_blas_threads(3) + for _ in range(4): + _limit_blas_threads(3) + assert _current_blas_threads() == 3 + + +def test_limit_is_process_wide_and_cross_thread(): + """BLAS pools are process-global, so a worker thread sees the cap too. + + This is the property that makes the fix work at all: the clustering call + runs on a request worker, not on the thread that built the model. + """ + _require_blas_pools() + + import threading + + _limit_blas_threads(2) + seen = [] + + def worker(): + seen.append( + sorted( + { + pool["num_threads"] + for pool in threadpoolctl.threadpool_info() + if pool["user_api"] == "blas" + } + ) + ) + + thread = threading.Thread(target=worker) + thread.start() + thread.join() + + assert seen == [[2]], "a worker thread did not observe the cap: %s" % seen + + +def test_the_cap_actually_reduces_blas_work(): + """The point of the change: fewer threads must mean less CPU. + + A 357-row Laplacian is the size a 35-minute recording produces. + """ + _require_blas_pools() + + import os + + clock = os.sysconf("SC_CLK_TCK") + + def process_cpu_seconds(): + fields = open("/proc/self/stat").read().split() + return (int(fields[13]) + int(fields[14])) / clock + + rng = np.random.RandomState(0) + laplacian = rng.rand(357, 357) + laplacian = laplacian + laplacian.T + + def measure(): + before = process_cpu_seconds() + scipy.linalg.eigh(laplacian) + return process_cpu_seconds() - before + + threads = os.cpu_count() or 1 + if threads < 4: + pytest.skip("needs at least 4 cores to show a difference") + + _limit_blas_threads(1) + one_thread_cpu = measure() + _limit_blas_threads(threads) + many_thread_cpu = measure() + + assert many_thread_cpu > one_thread_cpu, ( + "capping BLAS did not reduce CPU: %d threads=%.3f core-s, " + "1 thread=%.3f core-s" % (threads, many_thread_cpu, one_thread_cpu) + )