From 98dba1802d9ea1362cd9532ad716c52de1c879fe Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Fri, 11 Sep 2026 07:22:02 +0200 Subject: [PATCH] Implement NUMA topology detection and CPU budget capping This change introduces logic to detect NUMA topology, calculate per-node worker limits based on available CPU cores, and cap resource allocation to ensure worker counts respect the physical capacity of each node. --- src/Cargo.lock | 2 +- src/obikmer/Cargo.toml | 2 +- src/obisys/src/numa/runner.rs | 32 ++++++++------ src/obisys/src/numa/topology.rs | 77 +++++++++++++++++++++++---------- 4 files changed, 76 insertions(+), 37 deletions(-) diff --git a/src/Cargo.lock b/src/Cargo.lock index 20ddaff7..947c0a9e 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1480,7 +1480,7 @@ dependencies = [ [[package]] name = "obikmer" -version = "1.3.2" +version = "1.3.3" dependencies = [ "clap", "csv", diff --git a/src/obikmer/Cargo.toml b/src/obikmer/Cargo.toml index 84c8fcb5..78b45a85 100644 --- a/src/obikmer/Cargo.toml +++ b/src/obikmer/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "obikmer" -version = "1.3.2" +version = "1.3.3" edition = "2024" [[bin]] diff --git a/src/obisys/src/numa/runner.rs b/src/obisys/src/numa/runner.rs index 29c583b8..25d6929f 100644 --- a/src/obisys/src/numa/runner.rs +++ b/src/obisys/src/numa/runner.rs @@ -41,9 +41,10 @@ struct NodeConfig { /// 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. +/// each node's own worker cap, applied per node, so the pace of ramp-up +/// depends on that node's size rather than the node count — a +/// single-NUMA-node (UMA) machine ramps just as fast as an 8-node one, and a +/// `--cpu-max`-emptied node simply never ramps. /// /// # Termination /// @@ -63,22 +64,27 @@ impl PartitionRunner { } /// Detect topology and build. Always succeeds. + /// + /// Each node's worker count is its own (`--cpu-max`-budgeted) CPU count, + /// not a single value shared across nodes — `build()` can leave some + /// nodes with fewer cores than others (or none at all) once the process + /// CPU budget is smaller than the raw hardware topology, and spawning a + /// uniform worker count per node regardless would silently exceed that + /// budget on the emptied-out nodes. pub fn new() -> Self { let ns = build(); - let wpn = ns.workers_per_node(); debug!( - "PartitionRunner: {} node(s) × {} worker(s)/node max", + "PartitionRunner: {} node(s), {:?} core(s)/node", ns.pools.len(), - wpn, + ns.cpus_per_node.iter().map(Vec::len).collect::>(), ); let nodes = ns .pools .into_iter() .zip(ns.cpus_per_node) - .map(|(pool, cpu_ids)| NodeConfig { - pool, - cpu_ids, - max_workers: wpn, + .map(|(pool, cpu_ids)| { + let max_workers = cpu_ids.len(); + NodeConfig { pool, cpu_ids, max_workers } }) .collect(); Self { nodes } @@ -113,9 +119,9 @@ impl PartitionRunner { /// 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 + /// `(node's max_workers / INITIAL_DIVISOR).max(1)` are woken immediately + /// on every node, then `(node's max_workers / 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 diff --git a/src/obisys/src/numa/topology.rs b/src/obisys/src/numa/topology.rs index 0263e90c..806d6a81 100644 --- a/src/obisys/src/numa/topology.rs +++ b/src/obisys/src/numa/topology.rs @@ -19,22 +19,22 @@ pub struct NumaSetup { pub cpus_per_node: Vec>, } -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. +/// +/// Every node's CPU list — and therefore its Rayon pool's thread count — is +/// capped against [`crate::cpu_budget`] (the process-wide ceiling set by +/// `--cpu-max`, or the cgroup/host default when unset) via +/// [`cap_to_budget`]. Sizing the pool itself, not just the outer worker +/// count layered on top in [`super::runner::PartitionRunner`], matters +/// because callers query `rayon::current_num_threads()` from *inside* a +/// pool-installed closure to size further internal parallelism — that call +/// only sees the requested budget if the pool itself was built that small. #[cfg(feature = "numa")] pub fn build() -> NumaSetup { + let budget = crate::cpu_budget(); if let Ok(topology) = Topology::new() { - let nodes: Vec> = topology + let mut nodes: Vec> = topology .objects_with_type(ObjectType::NUMANode) .filter_map(|obj| obj.cpuset()) .map(|cpuset| { @@ -47,15 +47,23 @@ pub fn build() -> NumaSetup { .collect(); if nodes.len() > 1 { + cap_to_budget(&mut nodes, budget); if let Some(pools) = nodes .iter() - .map(|cpus| build_pool(cpus).map(|p| Some(Arc::new(p)))) + .map(|cpus| { + if cpus.is_empty() { + Some(None) + } else { + build_pool(cpus).map(|p| Some(Arc::new(p))) + } + }) .collect::>>() { debug!( - "NUMA topology: {} node(s), {} core(s)/node", + "NUMA topology: {} node(s), {} core(s)/node, budget {}", nodes.len(), nodes.first().map_or(0, |v| v.len()), + budget, ); return NumaSetup { pools, @@ -65,8 +73,17 @@ pub fn build() -> NumaSetup { } } - // UMA fallback: single synthetic node, all cores, no pool, no pinning. - let n_cores = crate::effective_parallelism(); + // UMA fallback: single synthetic node, budget-capped cores, no pool, no pinning. + debug!("UMA: single synthetic node, {} core(s)", budget); + NumaSetup { + pools: vec![None], + cpus_per_node: vec![(0..budget).collect()], + } +} + +#[cfg(not(feature = "numa"))] +pub fn build() -> NumaSetup { + let n_cores = crate::cpu_budget(); debug!("UMA: single synthetic node, {} core(s)", n_cores); NumaSetup { pools: vec![None], @@ -74,13 +91,29 @@ pub fn build() -> NumaSetup { } } -#[cfg(not(feature = "numa"))] -pub fn build() -> NumaSetup { - let n_cores = crate::effective_parallelism(); - debug!("UMA: single synthetic node, {} core(s)", n_cores); - NumaSetup { - pools: vec![None], - cpus_per_node: vec![(0..n_cores).collect()], +/// Trims each NUMA node's CPU list, in place, so the total across all nodes +/// never exceeds `budget` — floor-split evenly across nodes. No-op when the +/// topology already fits within `budget`. When `budget` is smaller than the +/// number of nodes, the trailing nodes are emptied entirely (one core each +/// to as many leading nodes as `budget` allows) rather than every node +/// keeping a token core that would collectively blow the budget. +#[cfg(feature = "numa")] +fn cap_to_budget(nodes: &mut [Vec], budget: usize) { + let total: usize = nodes.iter().map(Vec::len).sum(); + if budget >= total || nodes.is_empty() { + return; + } + let per_node = (budget / nodes.len()).max(1); + for cpus in nodes.iter_mut() { + cpus.truncate(per_node); + } + let mut used = 0; + for cpus in nodes.iter_mut() { + if used >= budget { + cpus.clear(); + } else { + used += cpus.len(); + } } }