Skip to main content

zsfm_tabfm/ensemble/
mod.rs

1pub mod aggregate;
2pub mod calibration;
3pub mod cat_encoder;
4pub mod config_gen;
5pub mod nnls;
6pub mod oof;
7pub mod orchestrate;
8pub mod pyrandom;
9pub mod scalers;
10
11/// Runs `f` inside a dedicated `rayon` thread pool of `threads` workers, or on rayon's own
12/// global default pool (all logical cores, or `RAYON_NUM_THREADS` if set) when `threads` is
13/// `None`. Shared by the CLI (`--threads`) and the Python bindings (`n_threads`) so both tune
14/// the same underlying `configs.par_iter()` loops in `orchestrate.rs`/`oof.rs`.
15///
16/// More threads than roughly the physical core count can *hurt* wall time here — Accelerate/BLAS
17/// does its own internal matmul threading, and the two can oversubscribe the machine. Benchmark
18/// before picking a non-default value for a given deployment.
19pub fn with_thread_pool<T>(
20    threads: Option<usize>,
21    f: impl FnOnce() -> anyhow::Result<T> + Send,
22) -> anyhow::Result<T>
23where
24    T: Send,
25{
26    match threads {
27        Some(n) => rayon::ThreadPoolBuilder::new()
28            .num_threads(n)
29            .build()?
30            .install(f),
31        None => f(),
32    }
33}