Eric Coissac
5f95e866f8
Introduce `obisys::effective_parallelism()` to read Linux cgroup v1/v2 CPU quotas from sysfs, preventing thread pool oversubscription in containerized environments. Replace direct `std::thread::available_parallelism()` calls across `obikindex` and `obikmer` with this centralized function. Bump `obikmer` version to 1.1.41.
494 lines
19 KiB
Rust
494 lines
19 KiB
Rust
// NUMA-aware partition runner via hwlocality.
|
||
//
|
||
// Detects NUMA topology using hwloc (cross-platform: Linux, macOS, etc.) and
|
||
// builds one Rayon ThreadPool per NUMA node with threads pinned to that node's
|
||
// CPUs. Linux first-touch policy then places graph allocations in local DRAM
|
||
// automatically — no explicit memory binding needed.
|
||
//
|
||
// UMA systems (single socket, Apple Silicon, etc.) are the degenerate case:
|
||
// one synthetic node containing all cores, no pool, no pinning.
|
||
|
||
use std::sync::Arc;
|
||
use std::time::{Duration, Instant};
|
||
|
||
use crossbeam_channel::unbounded;
|
||
#[cfg(feature = "numa")]
|
||
use hwlocality::Topology;
|
||
#[cfg(feature = "numa")]
|
||
use hwlocality::cpu::binding::CpuBindingFlags;
|
||
#[cfg(feature = "numa")]
|
||
use hwlocality::cpu::cpuset::CpuSet;
|
||
#[cfg(feature = "numa")]
|
||
use hwlocality::object::types::ObjectType;
|
||
use obisys::{CpuSample, IoSample};
|
||
use tracing::debug;
|
||
|
||
// ── Public interface ──────────────────────────────────────────────────────────
|
||
|
||
pub struct NumaSetup {
|
||
/// One entry per NUMA node. `None` on UMA systems (no pool, no pinning).
|
||
pub pools: Vec<Option<Arc<rayon::ThreadPool>>>,
|
||
/// CPU indices for each NUMA node, in node order.
|
||
pub cpus_per_node: Vec<Vec<usize>>,
|
||
}
|
||
|
||
impl NumaSetup {
|
||
/// Maximum worker slots per node (one per physical core in the node).
|
||
pub fn workers_per_node(&self) -> usize {
|
||
self.cpus_per_node
|
||
.first()
|
||
.map(|c| c.len().max(1))
|
||
.unwrap_or(1)
|
||
}
|
||
}
|
||
|
||
/// Detect NUMA topology and build per-node Rayon pools.
|
||
/// Always succeeds: falls back to a single synthetic UMA node on failure.
|
||
#[cfg(feature = "numa")]
|
||
pub fn build() -> NumaSetup {
|
||
if let Ok(topology) = Topology::new() {
|
||
let nodes: Vec<Vec<usize>> = topology
|
||
.objects_with_type(ObjectType::NUMANode)
|
||
.filter_map(|obj| obj.cpuset())
|
||
.map(|cpuset| {
|
||
cpuset
|
||
.iter_set()
|
||
.map(|idx| usize::from(idx))
|
||
.collect::<Vec<_>>()
|
||
})
|
||
.filter(|v| !v.is_empty())
|
||
.collect();
|
||
|
||
if nodes.len() > 1 {
|
||
if let Some(pools) = nodes
|
||
.iter()
|
||
.map(|cpus| build_pool(cpus).map(|p| Some(Arc::new(p))))
|
||
.collect::<Option<Vec<_>>>()
|
||
{
|
||
debug!(
|
||
"NUMA topology: {} node(s), {} core(s)/node",
|
||
nodes.len(),
|
||
nodes.first().map_or(0, |v| v.len()),
|
||
);
|
||
return NumaSetup {
|
||
pools,
|
||
cpus_per_node: nodes,
|
||
};
|
||
}
|
||
}
|
||
}
|
||
|
||
// UMA fallback: single synthetic node, all cores, no pool, no pinning.
|
||
let n_cores = obisys::effective_parallelism();
|
||
debug!("UMA: single synthetic node, {} core(s)", n_cores);
|
||
NumaSetup {
|
||
pools: vec![None],
|
||
cpus_per_node: vec![(0..n_cores).collect()],
|
||
}
|
||
}
|
||
|
||
#[cfg(not(feature = "numa"))]
|
||
pub fn build() -> NumaSetup {
|
||
let n_cores = obisys::effective_parallelism();
|
||
debug!("UMA: single synthetic node, {} core(s)", n_cores);
|
||
NumaSetup {
|
||
pools: vec![None],
|
||
cpus_per_node: vec![(0..n_cores).collect()],
|
||
}
|
||
}
|
||
|
||
/// Bind the calling thread to `cpu_indices` using hwloc.
|
||
/// Silently returns on any error so the thread still runs, just unbound.
|
||
#[cfg(feature = "numa")]
|
||
pub fn pin_current_thread(cpu_indices: &[usize]) {
|
||
let Ok(topology) = Topology::new() else {
|
||
return;
|
||
};
|
||
let mut cpuset = CpuSet::new();
|
||
for &idx in cpu_indices {
|
||
cpuset.set(idx);
|
||
}
|
||
let _ = topology.bind_cpu(&cpuset, CpuBindingFlags::THREAD);
|
||
}
|
||
|
||
#[cfg(not(feature = "numa"))]
|
||
pub fn pin_current_thread(_cpu_indices: &[usize]) {}
|
||
|
||
// ── Internal helpers ──────────────────────────────────────────────────────────
|
||
|
||
#[cfg(feature = "numa")]
|
||
fn build_pool(cpus: &[usize]) -> Option<rayon::ThreadPool> {
|
||
let cpus = cpus.to_vec();
|
||
rayon::ThreadPoolBuilder::new()
|
||
.num_threads(cpus.len())
|
||
.spawn_handler(move |thread| {
|
||
let cpus = cpus.clone();
|
||
std::thread::Builder::new().spawn(move || {
|
||
pin_current_thread(&cpus);
|
||
thread.run();
|
||
})?;
|
||
Ok(())
|
||
})
|
||
.build()
|
||
.ok()
|
||
}
|
||
|
||
// ── PartitionRunner ─────────────────────────────────────────────────────────
|
||
|
||
/// Growth step (fraction of a node's worker capacity added per activation
|
||
/// event, see [`NodeActivation::grow`]).
|
||
const GROWTH_DIVISOR: usize = 8;
|
||
/// Minimum CPU efficiency growth to activate more workers, as a fraction of
|
||
/// the size of the *last growth step* (e.g. `0.2` after adding 8 workers
|
||
/// requires the next check to show at least +1.6 cores of growth — 20 % of
|
||
/// the ~8 cores those 8 workers should contribute if the workload is truly
|
||
/// CPU-bound). Scaling by the last step's size — not the cumulative total —
|
||
/// keeps the bar meaningful regardless of how many workers are already
|
||
/// active, instead of demanding an ever-larger absolute jump as the pool
|
||
/// grows.
|
||
const CPU_SPAWN_THRESHOLD: f64 = 0.2;
|
||
/// Minimum I/O throughput growth (relative) to activate more workers.
|
||
const IO_SPAWN_THRESHOLD: f64 = 0.2;
|
||
|
||
struct NodeConfig {
|
||
pool: Option<Arc<rayon::ThreadPool>>,
|
||
cpu_ids: Vec<usize>,
|
||
max_workers: usize,
|
||
}
|
||
|
||
/// Generic NUMA-aware runner for partition-level parallel work.
|
||
///
|
||
/// Workers are distributed evenly across NUMA nodes and pinned to their
|
||
/// node's CPUs. UMA is the degenerate case: one node, no pinning.
|
||
///
|
||
/// Workers are pre-spawned dormant, one activation channel per node so
|
||
/// growth always targets a specific node rather than whichever dormant
|
||
/// worker happens to wake up first on a shared channel. Growth (both the
|
||
/// initial count and each subsequent step) is expressed as a fraction of
|
||
/// `workers_per_node`, applied identically to every node, so the pace of
|
||
/// ramp-up depends on node size rather than node count — a single-NUMA-node
|
||
/// (UMA) machine ramps just as fast as an 8-node one.
|
||
///
|
||
/// # Termination
|
||
///
|
||
/// ```text
|
||
/// drop(part_tx) → part_rx drains → workers exit → drop their result_tx
|
||
/// drop(result_tx) → result_rx closes → controller loop exits
|
||
/// drop(activate_txs) → dormant workers exit cleanly
|
||
/// ```
|
||
pub struct PartitionRunner {
|
||
nodes: Vec<NodeConfig>,
|
||
}
|
||
|
||
impl PartitionRunner {
|
||
/// Total worker slots across all nodes.
|
||
pub fn max_workers(&self) -> usize {
|
||
self.nodes.iter().map(|n| n.max_workers).sum()
|
||
}
|
||
|
||
/// Detect topology and build. Always succeeds.
|
||
pub fn new() -> Self {
|
||
let ns = build();
|
||
let wpn = ns.workers_per_node();
|
||
debug!(
|
||
"PartitionRunner: {} node(s) × {} worker(s)/node max",
|
||
ns.pools.len(),
|
||
wpn,
|
||
);
|
||
let nodes = ns
|
||
.pools
|
||
.into_iter()
|
||
.zip(ns.cpus_per_node)
|
||
.map(|(pool, cpu_ids)| NodeConfig {
|
||
pool,
|
||
cpu_ids,
|
||
max_workers: wpn,
|
||
})
|
||
.collect();
|
||
Self { nodes }
|
||
}
|
||
|
||
/// Run `f(i)` for every index in `order`.
|
||
///
|
||
/// Workers are pre-spawned dormant and activated adaptively, per node:
|
||
/// `(workers_per_node / INITIAL_DIVISOR).max(1)` are woken immediately on
|
||
/// every node, then `(workers_per_node / GROWTH_DIVISOR).max(1)` more per
|
||
/// node each time the check below fires. A timer thread fires that check
|
||
/// every `TIMER_SECS` seconds; each completed partition resets that timer
|
||
/// (forcing an immediate check) and also triggers its own inline check. A
|
||
/// growth step happens whenever CPU efficiency grows by at least
|
||
/// `CPU_SPAWN_THRESHOLD` of what the last growth step should have
|
||
/// contributed, or I/O throughput grows by at least `IO_SPAWN_THRESHOLD`
|
||
/// (relative) since the last check — whichever resource is the actual
|
||
/// bottleneck still shows headroom.
|
||
///
|
||
/// `on_done(i, result, elapsed)` is called from the controller thread as
|
||
/// each partition completes — suitable for progress bars and result
|
||
/// aggregation.
|
||
///
|
||
/// Returns the first error produced by `f`, if any.
|
||
pub fn run<F, R, E, C>(&self, order: &[usize], f: F, mut on_done: C) -> Result<(), E>
|
||
where
|
||
F: Fn(usize) -> Result<R, E> + Send + Sync,
|
||
R: Send,
|
||
E: Send,
|
||
C: FnMut(usize, R, Duration) + Send,
|
||
{
|
||
let n_total = order.len();
|
||
if n_total == 0 {
|
||
return Ok(());
|
||
}
|
||
|
||
const TIMER_SECS: u64 = 30;
|
||
const INITIAL_DIVISOR: usize = 4;
|
||
|
||
// ── Channels ──────────────────────────────────────────────────────────
|
||
let (part_tx, part_rx) = unbounded::<usize>();
|
||
// reset_tx: controller → timer ("reset the 30 s window")
|
||
let (reset_tx, reset_rx) = unbounded::<()>();
|
||
// event_tx: workers + timer → controller (unified event stream)
|
||
let (event_tx, event_rx) = unbounded::<WorkerEvent<R, E>>();
|
||
// One activation channel per node: growth always targets a specific
|
||
// node, rather than whichever dormant worker happens to win the race
|
||
// on a channel shared across all nodes.
|
||
let (activate_txs, activate_rxs): (Vec<_>, Vec<_>) =
|
||
(0..self.nodes.len()).map(|_| unbounded::<()>()).unzip();
|
||
|
||
for &i in order {
|
||
part_tx.send(i).ok();
|
||
}
|
||
drop(part_tx);
|
||
|
||
let max_workers = self.max_workers();
|
||
let node_caps: Vec<usize> = self.nodes.iter().map(|n| n.max_workers).collect();
|
||
let f = &f;
|
||
|
||
let mut first_err: Option<E> = None;
|
||
|
||
std::thread::scope(|s| {
|
||
// ── Timer thread ──────────────────────────────────────────────────
|
||
// Sends TimerTick every TIMER_SECS seconds. Resets its window each
|
||
// time reset_rx receives a message (i.e. on partition completion).
|
||
let timer_tx = event_tx.clone();
|
||
s.spawn(move || {
|
||
let period = Duration::from_secs(TIMER_SECS);
|
||
loop {
|
||
crossbeam_channel::select! {
|
||
recv(reset_rx) -> r => {
|
||
if r.is_err() { break; } // reset_tx dropped → exit
|
||
}
|
||
default(period) => {
|
||
if timer_tx.send(WorkerEvent::TimerTick).is_err() { break; }
|
||
}
|
||
}
|
||
}
|
||
});
|
||
|
||
// ── Pre-spawn workers dormant, grouped by node ────────────────────
|
||
// Each worker listens on its own node's activation channel only.
|
||
for (node, arx) in self.nodes.iter().zip(activate_rxs.iter()) {
|
||
let cpu_ids = &node.cpu_ids;
|
||
for _ in 0..node.max_workers {
|
||
let prx = part_rx.clone();
|
||
let etx = event_tx.clone();
|
||
let arx = arx.clone();
|
||
let pool = node.pool.clone();
|
||
|
||
s.spawn(move || {
|
||
if arx.recv().is_err() {
|
||
return;
|
||
}
|
||
if !cpu_ids.is_empty() {
|
||
pin_current_thread(cpu_ids);
|
||
}
|
||
for i in &prx {
|
||
let t = Instant::now();
|
||
let r = match &pool {
|
||
Some(p) => p.install(|| f(i)),
|
||
None => f(i),
|
||
};
|
||
etx.send(WorkerEvent::Completed(i, r, t.elapsed())).ok();
|
||
}
|
||
});
|
||
}
|
||
}
|
||
// Drop controller's event_tx: event_rx closes when all workers +
|
||
// timer have exited.
|
||
drop(event_tx);
|
||
|
||
// ── Controller ────────────────────────────────────────────────────
|
||
let mut activation = NodeActivation::new(&activate_txs, &node_caps, max_workers);
|
||
activation.activate_initial(INITIAL_DIVISOR, n_total);
|
||
|
||
let mut cpu_sample = CpuSample::now();
|
||
let mut io_sample = IoSample::now();
|
||
let mut completed = 0usize;
|
||
|
||
while completed < n_total {
|
||
let Ok(event) = event_rx.recv() else { break };
|
||
match event {
|
||
WorkerEvent::Completed(i, r, dur) => {
|
||
match r {
|
||
Ok(v) => on_done(i, v, dur),
|
||
Err(e) => {
|
||
if first_err.is_none() {
|
||
first_err = Some(e);
|
||
}
|
||
}
|
||
}
|
||
completed += 1;
|
||
// Reset the 30 s timer.
|
||
reset_tx.send(()).ok();
|
||
// Inline check: same logic as a timer tick.
|
||
maybe_activate(
|
||
&mut activation,
|
||
&mut cpu_sample,
|
||
&mut io_sample,
|
||
completed,
|
||
n_total,
|
||
);
|
||
}
|
||
WorkerEvent::TimerTick => {
|
||
maybe_activate(
|
||
&mut activation,
|
||
&mut cpu_sample,
|
||
&mut io_sample,
|
||
completed,
|
||
n_total,
|
||
);
|
||
}
|
||
}
|
||
}
|
||
|
||
// Dormant workers exit once every sender for their node's channel
|
||
// is dropped — `activate_txs` holds the only ones.
|
||
drop(activate_txs);
|
||
// Timer thread exits when reset_tx closes.
|
||
drop(reset_tx);
|
||
});
|
||
|
||
match first_err {
|
||
Some(e) => Err(e),
|
||
None => Ok(()),
|
||
}
|
||
}
|
||
}
|
||
|
||
// ── Internal event type ───────────────────────────────────────────────────────
|
||
|
||
enum WorkerEvent<R, E> {
|
||
Completed(usize, Result<R, E>, Duration),
|
||
TimerTick,
|
||
}
|
||
|
||
/// Tracks how many of each node's dormant workers have been woken, and
|
||
/// grows every node by the same amount at each step (capped by that node's
|
||
/// remaining dormant workers and by the run's total budget) so load stays
|
||
/// balanced across nodes at every point in time — never just "one more
|
||
/// worker somewhere". Also remembers the size of the last real growth step
|
||
/// (`last_step`), used to scale the CPU activation threshold to what that
|
||
/// step could plausibly have contributed (see `maybe_activate`).
|
||
struct NodeActivation<'a> {
|
||
txs: &'a [crossbeam_channel::Sender<()>],
|
||
caps: &'a [usize],
|
||
active: Vec<usize>,
|
||
total: usize,
|
||
max: usize,
|
||
last_step: usize,
|
||
}
|
||
|
||
impl<'a> NodeActivation<'a> {
|
||
fn new(txs: &'a [crossbeam_channel::Sender<()>], caps: &'a [usize], max: usize) -> Self {
|
||
Self {
|
||
txs,
|
||
caps,
|
||
active: vec![0; txs.len()],
|
||
total: 0,
|
||
max,
|
||
last_step: 0,
|
||
}
|
||
}
|
||
|
||
fn total(&self) -> usize {
|
||
self.total
|
||
}
|
||
fn last_step(&self) -> usize {
|
||
self.last_step
|
||
}
|
||
fn max(&self) -> usize {
|
||
self.max
|
||
}
|
||
fn is_full(&self) -> bool {
|
||
self.total >= self.max
|
||
}
|
||
|
||
/// Wake up to `(node_cap / divisor).max(1)` dormant workers on every
|
||
/// node, capped by `n_total`. Called once at startup, unconditionally.
|
||
fn activate_initial(&mut self, divisor: usize, n_total: usize) {
|
||
self.grow(divisor, n_total);
|
||
}
|
||
|
||
/// Same per-node sizing as [`activate_initial`](Self::activate_initial),
|
||
/// applied as a growth step. Returns the number of workers actually
|
||
/// activated (may be less than requested once a node or the total
|
||
/// budget is exhausted). Updates `last_step` when it actually grew.
|
||
fn grow(&mut self, divisor: usize, n_total: usize) -> usize {
|
||
let before = self.total;
|
||
for idx in 0..self.txs.len() {
|
||
let wanted = (self.caps[idx] / divisor).max(1);
|
||
let room = self.caps[idx].saturating_sub(self.active[idx]);
|
||
let grow = wanted.min(room).min(n_total.saturating_sub(self.total));
|
||
for _ in 0..grow {
|
||
self.txs[idx].send(()).ok();
|
||
}
|
||
self.active[idx] += grow;
|
||
self.total += grow;
|
||
}
|
||
let grew = self.total - before;
|
||
if grew > 0 {
|
||
self.last_step = grew;
|
||
}
|
||
grew
|
||
}
|
||
}
|
||
|
||
fn maybe_activate(
|
||
activation: &mut NodeActivation,
|
||
cpu_sample: &mut CpuSample,
|
||
io_sample: &mut IoSample,
|
||
completed: usize,
|
||
n_total: usize,
|
||
) {
|
||
if activation.is_full() || completed >= n_total {
|
||
return;
|
||
}
|
||
|
||
// Expect roughly 1 core of extra efficiency per worker activated in the
|
||
// last growth step (CPU-bound case); require at least CPU_SPAWN_THRESHOLD
|
||
// (20 %) of that expected gain before growing again. Scaling by the last
|
||
// step's size — not the cumulative total — keeps the bar meaningful
|
||
// regardless of how many workers are already active: growing by 8 should
|
||
// always take ~+1.6 cores to confirm, whether that's the 2nd growth step
|
||
// or the 20th.
|
||
let cpu_threshold = CPU_SPAWN_THRESHOLD * activation.last_step() as f64;
|
||
|
||
// Call both unconditionally (no `||` short-circuit): each sampler must
|
||
// advance its own window every tick, regardless of what the other one
|
||
// reports, or it would starve behind whichever signal fires first.
|
||
let cpu_wants_more = cpu_sample.do_i_activate(cpu_threshold);
|
||
let io_wants_more = io_sample.do_i_activate(IO_SPAWN_THRESHOLD * activation.last_step() as f64);
|
||
if !(cpu_wants_more || io_wants_more) {
|
||
return;
|
||
}
|
||
|
||
let grew = activation.grow(GROWTH_DIVISOR, n_total);
|
||
if grew > 0 {
|
||
debug!(
|
||
"activated {} worker(s) — {}/{} active",
|
||
grew,
|
||
activation.total(),
|
||
activation.max()
|
||
);
|
||
}
|
||
}
|