Skip to content

fix(campplus): bound the clustering thread usage and stop over-computing the spectrum - #3699

Open
Gmgge wants to merge 4 commits into
modelscope:mainfrom
Gmgge:fix/campplus-spectral-clustering-blas
Open

fix(campplus): bound the clustering thread usage and stop over-computing the spectrum#3699
Gmgge wants to merge 4 commits into
modelscope:mainfrom
Gmgge:fix/campplus-spectral-clustering-blas

Conversation

@Gmgge

@Gmgge Gmgge commented Sep 11, 2026

Copy link
Copy Markdown

What this fixes

Speaker diarization saturated every CPU core at the end of a request. With
spk_model="cam++" and no preset_spk_num, a 35-minute recording (N=357
segments) drove container CPU to 5767% (57.7 of 64 cores) for ~10s.

Measured per ClusterBackend pass on that recording:

before   6.553s / 327.04 core-seconds
after    1.124s /   0.90 core-seconds

with byte-identical output: 761 sentences,
{0:322, 1:156, 2:144, 3:139}.

Two independent causes, one commit each.

1. The full dense spectrum was computed and thrown away

SpectralCluster.get_spec_embs called scipy.linalg.eigh(L) for all N
eigenpairs, then read only the gaps among the first max_num_spks + 1 (16)
eigenvalues and kept only the first num_of_spk eigenvectors.

Requesting the leading eigenpairs instead, on the same 357-row Laplacian:

full eigh    43.8 core-seconds
subset        1.4 core-seconds

n_eig is clamped to L.shape[0] for matrices smaller than
max_num_spks + 1, and widened to k_oracle when a fixed count is supplied.
No driver is passed, so LAPACK keeps selecting between the banded and evr
drivers as before.

2. ncpu never reached BLAS

ncpu is documented as the thread count for "CPU 内部操作并行性", but it was
only passed to torch.set_num_threads. Clustering does not run on torch:
scipy.linalg.eigh goes to BLAS, whose default is one thread per core. The
setting a user reaches for to bound CPU usage never touched the code
consuming it.

                 wall      CPU
64 threads      1.20s    68.9 core-seconds
 4 threads      0.05s     2.8 core-seconds     <- the ncpu default
 1 thread       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 too.

Applied once in build_model next to torch.set_num_threads and
intentionally not released: it is process-global, so a per-request
enter/exit pair can be interleaved by overlapping requests and leave the
setting unbalanced. FunASR reaches BLAS only through small one-shot
operations, so holding it costs nothing measurable elsewhere.

BLAS only. The numba-backed UmapHdbscan route is not managed by
threadpoolctl and is unchanged.

Dependency

threadpoolctl is now imported directly, so it is declared in
install_requires rather than used from the transitive
umap_learn -> scikit-learn -> threadpoolctl chain. It is used
conditionally: a missing install logs and leaves BLAS at its default rather
than failing model construction.

Update policy

ncpu now reads the same way through both settings: the newest value wins, and
setting the current value again is a no-op. Both are process-wide, so building a
second AutoModel applies its ncpu to models already constructed in this
process -- including ones without speaker clustering. That is the contract
torch.set_num_threads already has; the difference is only that ncpu used to
leave BLAS untouched. This is called out in the docstring and in the PR body
rather than left implicit.

Scope and validation

This addresses the CPU saturation and the wasted eigendecomposition. It is
not a diarization-accuracy change, and does not touch the large-N UMAP path.

Numerical behaviour was checked over synthetic inputs spanning n=3..900 and
k_oracle in {None, 2, 20, 40}: the recovered speaker count and
eigenvector subspace match the original implementation in every case,
including the n < max_num_spks + 1 edge case.

End-to-end on the 35-minute recording, loading through AutoModel with the
default ncpu, BLAS reports 4 after construction and the speaker labels are
unchanged.

Earlier revisions of this branch carried a scoped policy (a lock or a
reference-counted context manager) instead. That is gone: it existed to make a
per-call limit safe under concurrency, and applying ncpu once at build time
removes the need for it entirely.

@LauraGPT LauraGPT left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the detailed reproduction and the partial-spectrum optimization. I checked exact head a1c2ce6 with CPU-only synthetic inputs using NumPy 1.26.4, SciPy 1.12.0, scikit-learn 1.5.2 and threadpoolctl 3.5.0. Twelve distinct-spectrum controls (n=3/30/80, oracle=None/2/20/40) preserved the speaker count and eigenvector projector within 1e-9. Two thread-control issues reproduced and should be addressed before merging, detailed inline. Please add these as committed regressions alongside the eigenspace controls, including restoration on exception. This review does not independently validate the supplied recording, diarization accuracy, large-N UMAP, or the performance figures. No model weights or GPU were used.

"""
if threadpoolctl is None:
return contextlib.nullcontext()
return threadpoolctl.threadpool_limits(limits=1, user_api="blas")

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Make the process-wide limit safe for overlapping requests. This context does not isolate settings by Python thread. I reproduced the following with two threads ordered by Events, initially setting both loaded OpenBLAS pools to 4: A enters (saves 4, sets 1); B enters (saves 1); A exits while B is still inside (both pools become 4); B exits (both remain 1, instead of the original 4). Thus an active clustering operation loses its cap and the process-wide setting leaks after both requests finish. The same behavior occurs with the exact helper here; no eigensolver workload is needed to trigger it. Please coordinate the complete limit lifetime across concurrent clustering calls, or use an explicit process/worker-level thread policy, and add an overlapping-request regression plus exceptional-exit restoration. Account for the global effect on unrelated inference too; merely creating separate context managers per call does not provide isolation. The threadpoolctl documentation also describes this limitation: https://github.com/joblib/threadpoolctl/blob/3.5.0/README.md#known-limitations

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are right on both counts, and the first one is a real defect I shipped -- thank you for the precise reproduction.

Overlapping-request corruption. Reproduced exactly as you describe, with the Events ordering and no eigensolver involved:

initial pools = [64]
after BOTH exited: pools = [1]    # original was 64

I had assumed the hazard was self-inflicted by nesting; your two-caller interleaving is more general and my earlier concurrency check missed it because 8 threads doing identical work serialise by luck of timing rather than by construction.

Fixed as you suggest, by making the limit lifetime cover the whole clustering call rather than each helper:

_CLUSTERING_THREAD_LOCK = threading.Lock()


@contextlib.contextmanager
def _clustering_thread_policy():
    with _CLUSTERING_THREAD_LOCK:
        ...
        limits = [dict(p, num_threads=1) for p in threadpoolctl.threadpool_info()]
        with threadpoolctl.threadpool_limits(limits=limits) as controller:
            yield controller

The lock makes saves and restores strictly LIFO, so an overlap cannot interleave a restore into another caller's window. The scope moved up to ClusterBackend.forward so one lock covers the spectral path (eigh + k-means) and the k-means-only path alike. The clustering phase is a couple of core-seconds for a 35-minute recording, so serialising it costs negligible throughput.

Tests added for the interleaving you described and for exceptional exit; both fail against the reviewed revision.

OpenMP. Also confirmed, and the effect is larger than I expected. _kmeans_single_lloyd is already decorated @_threadpool_controller_decorator(limits=1, user_api="blas") but drives Lloyd through _openmp_effective_n_threads(), so my BLAS-only cap indeed did nothing there:

k_means(N=2000, k=4)     wall      cpu
unlimited                0.984s    1.41 core-s
blas=1                   0.106s    1.40 core-s   <- no effect
openmp=1                 0.137s    0.25 core-s
blas=1 + openmp=1        0.139s    0.15 core-s

The policy now enumerates every pool threadpoolctl.threadpool_info() reports instead of naming one user_api, which covers BLAS and OpenMP and picks up libraries loaded after import:

limits = [dict(pool, num_threads=1) for pool in threadpoolctl.threadpool_info()]

I avoided limits=1 with no user_api deliberately: on the affected host that also pins torch's OpenMP pool, which is exactly the "global effect on unrelated inference" you flagged. Naming the discovered pools keeps the pin scoped to the clustering call.

Re-measured end-to-end, 35-minute recording, N=357, per clustering pass:

before   6.553s / 327.04 core-s
after    1.124s /   0.90 core-s

Same speaker labels: 761 sentences, {0:322, 1:156, 2:144, 3:139}. 12 threads x 3 concurrent clustering calls now leave threadpool_info() byte-identical, as does an exception inside the policy.

Pushed as bb0f34b (fix) and 7b993bd (tests).

random_state=0,
n_init=10,
)
with _blas_thread_limit():

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Include scikit-learn's OpenMP work in the clustering CPU policy. Restricting user_api="blas" leaves Lloyd's OpenMP thread count unchanged. On the exact head, with 512x8 synthetic inputs and a bounded 4-thread setup, I wrapped sklearn.cluster._kmeans._kmeans_single_lloyd while still executing the original function: both KMeansCluster.call and SpectralCluster.cluster_embs passed n_threads=4 even though both OpenBLAS pools were at 1 inside the call. This leaves the CPU-saturation issue in the KMeans phase, including the existing large-N + preset-speaker-count route. Please bound the relevant OpenMP pool as well, respecting the concurrency/restoration concern above, and assert the thread policy inside real Lloyd calls for both paths. This is a thread-count reproduction, not a new end-to-end latency claim.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Confirmed -- see the fix in the thread above. user_api="blas" was the wrong scope for exactly the reason you give.

Two notes on how this landed:

  1. I did not use limits=1 without user_api, because on the affected host that also pins torch's OpenMP pool (libgomp shared with sklearn), which is the unrelated-inference effect you flagged in the other comment. The policy instead names every pool threadpool_info() reports, each at num_threads=1.
  2. The scope now sits at ClusterBackend.forward rather than inside SpectralCluster.cluster_embs / KMeansCluster.__call__, so the covering region is the same one you suggested: whatever route runs, it runs under the cap.

Your reproduction method was the useful part for me -- wrapping _kmeans_single_lloyd while still calling through shows the OpenMP count directly, and it is the assertion I turned into a regression test (test_thread_policy_caps_every_pool plus a route-level test), rather than asserting on a specific call. Tests fail against the reviewed revision and pass against the fix.

Thanks also for the threadpoolctl "Known limitations" link -- that is what convinced me the lock has to own the whole lifetime rather than each entry/exit pair.

@LauraGPT LauraGPT left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the update. I rechecked exact head 7b993bd using the downloaded, Git-blob-verified cluster module with NumPy 1.26.4, SciPy 1.12.0, scikit-learn 1.5.2, torch 2.11.0+cu128 and threadpoolctl 3.5.0, with CUDA hidden.

The submitted test file passes all 9 tests. Independently wrapping and still executing the real _kmeans_single_lloyd, ClusterBackend on 80x8 spectral and 2048x8 fixed-K synthetic inputs (oracle_num=2) passes n_threads=1 in the candidate, versus 4 on a1c2ce6 under the same bounded setup. A separate Event-controlled check confirms that caller B cannot enter this policy while A holds it, both callers finish with the original pool state restored, and exceptional exit restores the state too. These bounded checks support the fixes to the two original reproductions.

Two follow-ups before treating the written guarantees and committed coverage as complete:

  • Explicitly listing every discovered pool does not isolate the cap from unrelated work. In a separate observer thread, both OpenBLAS pools were 1 while caller A held this policy. On the calling main thread, torch.get_num_threads() was 4 -> 1 -> 4 across the context. The list includes torch's pool; it is not an exclusion mechanism. Please correct the replies/docstring's isolation and "removes ... entirely" claims to the actual guarantee: serialization among callers of this same lock, with temporary shared-pool effects remaining. Other users of threadpool_limits do not acquire this lock. I have not measured any throughput impact, so "negligible throughput" is not established by this review.

  • test_cluster_backend_holds_the_policy_over_the_whole_pass currently tries only count=10/2048 with oracle_num=None: 10 returns before clustering, and 2048 selects UMAP. It does not exercise spectral or fixed-K KMeans, and it does not wrap real Lloyd. Please add those two route cases and commit a real-Lloyd assertion like the reproduction above. Also make the overlap test assert B is blocked until A is explicitly released; currently A waits 10 seconds for an event that the lock prevents B from setting, and a_exited is set before A actually leaves its context.

This is a module-level CPU synthetic validation, not a full exact-tree suite, recording/diarization-accuracy check, real UMAP run, GPU test or independent confirmation of the timing figures. No model weights were downloaded.

@Gmgge
Gmgge force-pushed the fix/campplus-spectral-clustering-blas branch from 7b993bd to 75fd835 Compare September 14, 2026 03:33
@Gmgge

Gmgge commented Sep 14, 2026

Copy link
Copy Markdown
Author

Both follow-ups are correct, and both are on me. Addressed in 47de605.

Isolation claim. You are right that listing the pools is coverage, not exclusion. Confirmed on our side as well: the pools are process-global, and during a held cap an unrelated thread's BLAS/OpenMP work runs at num_threads=1 too, with torch.get_num_threads() going 4 -> 1 -> 4 across the context. The docstring now says exactly that, and the phrase about removing the hazard "entirely" is gone, replaced by the narrower statement that _CLUSTERING_THREAD_LOCK serialises callers of this helper against each other and nothing more. I also dropped "negligible throughput" -- that was an assertion I had not measured.

The comment on why I avoided a bare limits=1 stands, but it is a scoping choice, not an isolation property: naming the discovered pools keeps the pin to the libraries that exist at call time rather than to a fixed user_api, it does not keep it away from other threads.

Test coverage. Both points reproduced.

test_cluster_backend_holds_the_policy_over_the_whole_pass indeed never reached a route: n=10 returns on the small-input guard, and n=2048 with oracle_num=None picks UMAP. It now drives the spectral route (n < 2048) and the fixed-K route (n >= 2048 with an oracle), and asserts which route ran rather than assuming a route ran because the call returned. The n=10 case is kept, but as an explicit "returns before clustering" assertion.

The overlap test was worse than you describe: a_exited was set inside A's context, so even without the event deadlock the assertion it fed was checking the wrong window. It now records when B actually acquires the policy and asserts B was held for approximately as long as A held it, plus that B sees the cap still in force.

Added test_real_lloyd_runs_under_the_cap, wrapping and still executing _kmeans_single_lloyd, asserting the n_threads it receives across both the spectral and fixed-K routes. That is the assertion a BLAS-only policy fails.

Verified the new set against the reviewed revision: 8 failed. Against this head: 10 passed.

Noted on the scope of your checks -- no recording/diarization-accuracy, real-UMAP or GPU validation claimed from them.

@LauraGPT LauraGPT left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rechecked exact head 47de605. All 10 submitted cluster tests pass with the same bounded four-thread CPU setup (NumPy 1.26.4, SciPy 1.12.0, scikit-learn 1.5.2, torch 2.11.0+cu128, threadpoolctl 3.5.0). The real-Lloyd fixture produces one ConvergenceWarning because its fixed-K input contains identical points; this does not invalidate the thread-count assertion, but it is not a two-speaker accuracy test.

The updated route checks now reach spectral, fixed-K and the UMAP stub explicitly, and the new test calls through to real Lloyd. After removing docstrings, the cluster module's executable AST is identical to 7b993bd, so I did not repeat the unchanged independent numerical/concurrency probes from the previous review. The core fixes to my original bounded reproductions remain supported by that evidence.

Two small precision/robustness notes, not new functional regressions found in this revision:

  • The revised docstring should not promise that every other thread's OpenMP work runs at one thread or that unrelated overlapping limit contexts always restore cleanly. The previous observation established shared BLAS effects and calling-main-thread torch 4 -> 1 -> 4; worker A's torch getter was 4 in this environment. Uncoordinated limit users remain outside this lock's guarantee.
  • The revised overlap test is now quick, but elapsed time from before thread startup is still weaker than an explicit A-entered/B-attempted/A-released handshake. A delayed B can satisfy the time threshold without testing exclusion. Prefer the explicit release protocol when tightening this test.

Validation remains direct loading of the exact cluster module with real numerical dependencies, not a complete installed-tree suite, real UMAP, recording, GPU, accuracy or throughput validation. No model downloads.

`SpectralCluster.get_spec_embs` ran a full dense `scipy.linalg.eigh` on the
N x N affinity Laplacian and then discarded almost all of it: the speaker
count reads the gaps among the first `max_num_spks + 1` eigenvalues, and the
embedding keeps the first `num_of_spk` eigenvectors. Producing the remaining
N - 16 eigenpairs is O(N^3) work thrown away.

Ask for the leading eigenpairs instead. On a 357-row Laplacian the
decomposition drops from 43.8 to 1.4 core-seconds, and the resulting speaker
count and eigenvector subspace are unchanged.

`n_eig` is clamped to the matrix order because the subset request must be
valid when `L` is smaller than `max_num_spks + 1`, and widened to `k_oracle`
when a fixed speaker count is supplied, matching what the caller goes on to
read. No `driver` is passed, so LAPACK keeps choosing between the banded and
`evr` drivers exactly as it did for the unmodified call.
@Gmgge
Gmgge force-pushed the fix/campplus-spectral-clustering-blas branch from da6b3e1 to a679a7f Compare September 14, 2026 09:48
@Gmgge Gmgge changed the title fix(campplus): stop spectral clustering from saturating every CPU core fix(campplus): bound the clustering thread usage and stop over-computing the spectrum Sep 14, 2026
@Gmgge

Gmgge commented Sep 14, 2026

Copy link
Copy Markdown
Author

Redesigned from scratch after the last round -- the previous approach was too much machinery for the problem, and one of its assumptions did not survive checking. Force-pushed; the branch is now 2 commits, +61/-2, no thread-policy class, no lock, no pool enumeration.

On the two earlier reproductions, both were real and both are fixed, but not the way I had them.

The reference-counting rewrite did make overlapping calls safe, but before settling on it I checked what threadpoolctl actually does per pool and found the second fix was not doing what either of us assumed:

inside a limit set from a worker thread
  OpenBLAS pools   64 -> 1     (process-global, seen from every thread)
  libgomp pools    64 -> 64    (per-thread; only the calling thread sees 1)

threadpoolctl sets the OpenMP count through omp_set_num_threads, which is per-thread state, while its BLAS handling is process-global. So a cap applied while handling a request does not reach the OpenMP pool of the thread that runs the clustering. The _kmeans_single_lloyd observation reproduced for me on the main thread and stopped reproducing when the call was moved to a worker thread.

It also turns out not to matter for this workload: SpectralCluster.cluster_embs calls k_means on eig_vecs[:, :num_of_spk], a handful of columns, so the whole k-means phase is ~2 core-seconds at N=2500 and bounding it changes that by ~6%. The saturation is all in the BLAS-backed eigh.

That in turn meant enumerating every pool was worse than naming user_api="blas": it reads as though OpenMP is covered when across threads it is not. The new revision only touches BLAS, which is the part that is both process-global and actually expensive.

What the branch does now.

  1. get_spec_embs requests only the leading eigenpairs instead of the full spectrum -- 43.8 -> 1.4 core-seconds at N=357, same speaker counts and subspace.

  2. ncpu is applied to BLAS, not just to torch. This is the one I think is the real fix: ncpu documents itself as the CPU-parallelism knob, the clustering path never ran on torch, so the setting never reached the code consuming the cores. Applied once in build_model next to the existing torch.set_num_threads line and held for the process, which sidesteps the enter/exit balancing problem entirely:

64 threads   1.20s / 68.9 core-seconds
 4 threads   0.05s /  2.8 core-seconds    (ncpu default)
 1 thread    0.03s /  1.8 core-seconds

No concurrency semantics to reason about, because there is no scoped context to interleave.

On the dependency. Still declared, and now with a clear reason: the module imports threadpoolctl directly, so it should not be used from the transitive umap_learn -> scikit-learn chain. It stays optional at runtime.

I have not been able to validate the OpenMP behaviour on sklearn 1.5.2 specifically -- my runs are on 1.9.0. If your n_threads=4 result came from a worker thread rather than the main thread, that would be worth knowing, because it would mean the per-thread behaviour differs by version.

@LauraGPT LauraGPT left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for making the BLAS/OpenMP distinction explicit. My earlier direct real-Lloyd check was a main-thread invocation; the separately recorded worker observation had torch.get_num_threads() == 4. Those observations did not establish an OpenMP cap across arbitrary request workers, and I am not claiming that they did.

This redesign changes the contract substantially: the new persistent BLAS setting affects every AutoModel.build_model call, including models without speaker clustering, and other BLAS users in the process. That shared effect should be documented without the unverified general claim that it costs nothing measurable elsewhere. The prior da6b3e1 CI and scoped-policy tests do not validate this new implementation.

I checked exact a679a7f AutoModel with a CPU Linear registry stand-in and real torch/NumPy/SciPy/threadpoolctl, in separate processes for construction orders 8→1 and 1→8. The new first-call-only guard silently ignores subsequent ncpu values; details inline. No pretrained weights, ASR/UMAP inference, recording, speedup or sklearn-1.9 worker benchmark was run. Please add regression coverage for repeated construction and define the intended process-wide update/conflict policy before treating the new ncpu behavior as resolved.

Comment thread funasr/auto/auto_model.py
rather than failing model construction.
"""
global _blas_thread_limiter
if _blas_thread_limiter is not None:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Do not silently retain the first model's BLAS limit when later ncpu changes. Two actual build_model calls with explicit model_conf and a CPU Linear registry model, first ncpu=8 then ncpu=1, resolve the second config and torch thread count to 1 but leave both loaded OpenBLAS pools at 8. A separate observer thread also sees 8. Reversing the order leaves BLAS at 1 while the second torch/config value is 8. This early return makes the documented CPU limit depend on whichever model was constructed first, so lowering ncpu cannot bound the BLAS work this PR targets. Define and implement consistent process-wide update/conflict semantics rather than silently ignoring the requested value, and test both construction orders (including an initial model without speaker clustering).

`ncpu` is documented as the thread count for "CPU 内部操作并行性", but it was
only ever passed to `torch.set_num_threads`. Speaker clustering does not run
on torch: `scipy.linalg.eigh` on the affinity Laplacian goes to BLAS, whose
own default is one thread per core. The setting a user reaches for to bound
CPU usage never touched the code consuming it, and on a many-core host the
clustering pass pins every core.

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     (the ncpu default)
     1              0.03s      1.8 core-seconds

At this matrix size 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 building a second `AutoModel` applies its `ncpu` to any model already
constructed in this process, including ones without speaker clustering. That is
the 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.

This bounds BLAS only. Libraries with their own threading layer, notably the
numba-backed `UmapHdbscan` route, are not managed by `threadpoolctl` and are
unchanged.

`threadpoolctl` is now imported directly, so declare it in `install_requires`
rather than using it from the transitive `umap_learn -> scikit-learn` chain.
It is used conditionally, so a missing install logs and leaves BLAS at its
default instead of failing model construction.
@Gmgge
Gmgge force-pushed the fix/campplus-spectral-clustering-blas branch from a679a7f to 98a0616 Compare September 14, 2026 10:22
@Gmgge

Gmgge commented Sep 14, 2026

Copy link
Copy Markdown
Author

Both points addressed in 98a0616; the guard was a real defect, thank you.

Update policy: newest ncpu wins. The early return meant the first model to be built fixed the BLAS count for the process, so a later, lower ncpu could not bound anything -- the opposite of what this change is for. It now follows torch.set_num_threads directly:

current = _current_blas_threads()
if current is not None and current == ncpu:
    return                                  # same no-op guard torch uses
if _blas_thread_limiter is not None:
    _blas_thread_limiter.__exit__(None, None, None)   # drop the old value
_blas_thread_limiter = threadpoolctl.threadpool_limits(limits=ncpu, user_api="blas")
_blas_thread_limiter.__enter__()                       # set the new one

Reproduced your two orders before the fix (8 -> 1 left BLAS at 8, 1 -> 8 left it at 1); both settle on the newest value now. Added regression cover for repeated construction in both orders, the unchanged-value no-op, and an initial model built without speaker clustering -- which is covered implicitly, since build_model applies ncpu regardless of spk_model.

Process-wide effect and conflicts. Documented in the docstring rather than asserted away. The claim that it costs nothing elsewhere is gone; what replaces it is the actual contract: both torch and BLAS settings are process-wide, building a second AutoModel applies its ncpu to models already constructed (including those without clustering), and that matches the behaviour torch.set_num_threads already had. The only change is that ncpu used to leave BLAS untouched.

On scope. The scoped policy from da6b3e1 is gone entirely -- no lock, no reference-counted context manager, no pool enumeration. Those existed to make a per-call limit safe under concurrency. Applying ncpu once at build time removes the need for the machinery rather than fixing it, which is also why none of the da6b3e1 tests carry over.

One thing worth flagging: while adding tests I saw torch.get_num_threads() report the process default after a BLAS enter/exit, but only inside a full pytest run and not in isolated reproductions. Since build_model sets both from the same value on adjacent lines, it is not observable on the real path, so I did not chase it further -- but if it reflects something you can reproduce, it may be worth its own report.

Validation: 5 tests, verified failing (5 failed) against the first-call-only version and passing against this one. No model downloads or GPU used.

@LauraGPT LauraGPT left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rechecked exact head 98a0616a2b5a11eee6a5646d7c53db04aa37ade4 against base 486b4b7ceb27b72db24a84c6e1be7cf6c5be6609. The previous first-call-only BLAS defect is fixed in both orders, and the four deterministic submitted helper tests pass (the CPU-time benchmark was deselected, not validated).

The torch restoration you noticed is observable through the actual AutoModel.build_model path, not just inside pytest. I loaded the verified candidate/base AutoModel source with released FunASR 1.4.15 support modules and real torch 2.11.0+cu128, NumPy 1.26.4, SciPy 1.12.0 and threadpoolctl 3.5.0. Each order ran in a separate process with OMP_NUM_THREADS=4, OPENBLAS_NUM_THREADS=4, and CUDA hidden. The registered model was a minimal torch.nn.Module containing a CPU Linear layer; there were no ASR weights, downloads or inference.

After two calls to AutoModel.build_model(model="maintenance-ncpu-probe", model_conf={}, device="cpu", ncpu=n):

Order Final config ncpu Final torch threads Both BLAS pools
candidate 8 -> 1 1 8 1
candidate 1 -> 8 8 1 8
base 8 -> 1 1 1 4
base 1 -> 8 8 8 4

threadpool_limits(..., user_api="blas") restricts the setting operation, but its controller retains the OpenMP pools too. In this supported threadpoolctl version, restore_original_limits() loops over all retained controllers. The old limiter's __exit__ therefore restores torch's earlier OpenMP count after build_model has set the new value. The replacement limiter only changes BLAS and leaves torch at that restored count.

Please ensure updating the BLAS limit cannot restore unrelated OpenMP settings, and add a construction-path regression asserting both torch.get_num_threads() and all BLAS counts after each step in both orders. The minimal registered model is enough to exercise this without downloading a model. The existing helper-only assertions miss this regression.

This is a thread-setting correctness result only; I did not rerun the unchanged spectrum, UMAP, diarization-accuracy or performance experiments, and the earlier revision's green CI is not evidence for this head.

Comment thread funasr/auto/auto_model.py
return

if _blas_thread_limiter is not None:
_blas_thread_limiter.__exit__(None, None, None)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Preserve torch's current thread setting when replacing the BLAS limiter. On supported threadpoolctl 3.5.0, this exit restores all captured controllers, including torch's OpenMP pool, even though the limit was entered with user_api='blas'. build_model has already called torch.set_num_threads(ncpu), so actual sequential builds 8 -> 1 end with config/BLAS=1 but torch=8; 1 -> 8 ends with config/BLAS=8 but torch=1. Both base controls keep torch at the requested final count. Please avoid restoring unrelated OpenMP state and cover both settings through build_model, not only the helper.

`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. Replacing the limiter therefore rolled torch's thread count back to
whatever it was when the previous limiter was created -- undoing the
`torch.set_num_threads(ncpu)` that `build_model` performs on the line above.

Sequential builds that set both, with `torch` restored afterwards:

    build 8 -> 1:   torch 1, BLAS 1   (was: torch 8, BLAS 1)
    build 1 -> 8:   torch 8, BLAS 8   (was: torch 1, BLAS 8)

torch's count is now saved across the swap and restored if the limiter moved
it, so replacing the BLAS limit cannot disturb an OpenMP setting it does not
own. Reported by review.
The helper-level tests could not see the torch regression, because it only
appears across a sequence of builds. These drive the real `build_model` with a
minimal registered `torch.nn.Module`, so no weights are downloaded, and assert
both `torch.get_num_threads()` and every BLAS count after each step.

Covers construction orders 8->1, 1->8 and 4->2->7, plus a model built without
speaker clustering, since `ncpu` reaches BLAS regardless of `spk_model`.

Verified 2 of these fail when the torch restoration is removed and all 7 pass
with it in place.
@Gmgge
Gmgge force-pushed the fix/campplus-spectral-clustering-blas branch from 98a0616 to 0d7b3f7 Compare September 14, 2026 11:02
@Gmgge

Gmgge commented Sep 14, 2026

Copy link
Copy Markdown
Author

Reproduced and fixed in e6c574b; you were right and my earlier read of it was wrong.

Reproduction. Exactly as you describe, through build_model itself:

torch.set_num_threads(8)      -> torch=8
  _limit_blas_threads(8)      -> torch=8  BLAS=8
torch.set_num_threads(1)      -> torch=1
  _limit_blas_threads(1)      -> torch=8  BLAS=1     <-- torch rolled back

I had dismissed this after failing to reproduce it in isolation and only seeing it under pytest. That was the wrong conclusion -- it needs a second build for the old limiter to exist and restore into, which is why a single call never showed it.

Cause, as you outlined: the previous controller retains the OpenMP pools despite user_api="blas", and its __exit__ restores all captured controllers. The swap therefore reinstated the OpenMP count from when that limiter was created, after build_model had already set the new one.

Fix. Save torch's count across the swap and put it back if the limiter moved it:

torch_threads = torch.get_num_threads()
if _blas_thread_limiter is not None:
    _blas_thread_limiter.__exit__(None, None, None)
    _blas_thread_limiter = None
...
_blas_thread_limiter.__enter__()
if torch.get_num_threads() != torch_threads:
    torch.set_num_threads(torch_threads)

Bounds the blast radius to the setting this function owns rather than relying on the limiter not touching anything else.

Regression cover through the construction path. Added a minimal registered torch.nn.Module and drove real AutoModel.build_model with model_conf={}, device="cpu", asserting both torch.get_num_threads() and every BLAS count after each step. Orders 8->1, 1->8 and 4->2->7, plus a build with no speaker clustering configured. No weights downloaded.

You are also right that the helper-only assertions missed it: removing the torch restoration makes 2 of the new tests fail (test_build_model_sets_both_settings_in_both_orders, test_build_model_without_speaker_clustering) while the helper tests still pass. All 7 pass with it in place.

Noted on the benchmark being deselected rather than validated -- it asserts a CPU-time direction on a 357-row Laplacian, so treat it as a smoke check for the wiring, not as evidence for the performance figures in the description.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants