Merge pull request 'Implement NUMA topology detection and CPU budget capping' (#74) from push-tnkxqmwztnmm into main

Reviewed-on: #74

src/Cargo.lock

@@ -1480,7 +1480,7 @@ dependencies = [
 
 [[package]]
 name = "obikmer"
-version = "1.3.2"
+version = "1.3.3"
 dependencies = [
  "clap",
  "csv",

src/obikmer/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "obikmer"
-version = "1.3.2"
+version = "1.3.3"
 edition = "2024"
 
 [[bin]]

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::<Vec<_>>(),
         );
         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

src/obisys/src/numa/topology.rs

@@ -19,22 +19,22 @@ pub struct NumaSetup {
     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.
+///
+/// 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<Vec<usize>> = topology
+        let mut nodes: Vec<Vec<usize>> = 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::<Option<Vec<_>>>()
             {
                 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<usize>], 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();
+        }
     }
 }