Feat: Introduce resource profiling and monitoring APIs

This change introduces a comprehensive profiling system, allowing for periodic resource usage monitoring (CPU, I/O, memory) via a background thread. It also exposes new public APIs for budget management, locking, NUMA partitioning, and progress tracking.

src/Cargo.lock

@@ -1480,7 +1480,7 @@ dependencies = [
 
 [[package]]
 name = "obikmer"
-version = "1.3.1"
+version = "1.3.4"
 dependencies = [
  "clap",
  "csv",
@@ -1504,6 +1504,7 @@ dependencies = [
  "obiread",
  "obiskbuilder",
  "obisys",
+ "rayon",
  "serde",
  "serde_yaml",
  "tracing",

src/obikmer/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "obikmer"
-version = "1.3.1"
+version = "1.3.4"
 edition = "2024"
 
 [[bin]]
@@ -28,6 +28,7 @@ obikrope      = { path = "../obikrope" }
 obifastwrite  = { path = "../obifastwrite" }
 obiskbuilder  = { path = "../obiskbuilder" }
 clap          = { version = "4", features = ["derive"] }
+rayon         = "1"
 csv           = "1"
 ndarray       = "0.17"
 serde         = { version = "1", features = ["derive"] }

src/obikmer/src/cli.rs

@@ -34,7 +34,9 @@ pub struct CommonArgs {
     #[arg(short, long, default_value_t = 256)]
     pub partitions: usize,
 
-    /// Number of worker threads
+    /// Number of worker threads. Silently clamped to the process CPU budget
+    /// (see the global `--cpu-max`); use it to run with *fewer* threads than
+    /// the budget, never more.
     #[arg(
         short = 'T',
         long,
@@ -85,9 +87,16 @@ impl CommonArgs {
         }
     }
 
+    /// Worker-thread count actually used: `--threads` clamped to the
+    /// process-wide CPU budget (`obisys::cpu_budget`, itself bounded by the
+    /// global `--cpu-max`).
+    pub fn effective_threads(&self) -> usize {
+        self.threads.min(obisys::cpu_budget()).max(1)
+    }
+
     pub fn effective_max_open(&self) -> usize {
         self.max_open_files
-            .unwrap_or_else(|| (self.threads / 4).max(1))
+            .unwrap_or_else(|| (self.effective_threads() / 4).max(1))
             .max(1)
     }
 

src/obikmer/src/cmd/index/mod.rs

@@ -243,7 +243,7 @@ pub fn run(args: IndexArgs) {
 
     // ── Stage 1: scatter ─────────────────────────────────────────────────────
     if current_state(&idx) < IndexState::Scattered {
-        let n_workers = args.common.threads.max(1);
+        let n_workers = args.common.effective_threads();
         let max_open = args.common.effective_max_open();
 
         let t = Stage::start("scatter");

src/obikmer/src/cmd/query/mod.rs

@@ -63,7 +63,9 @@ pub struct QueryArgs {
     #[arg(short = 'z', long)]
     pub findere_z: Option<usize>,
 
-    /// Number of worker threads
+    /// Number of worker threads. Silently clamped to the process CPU budget
+    /// (see the global `--cpu-max`); use it to run with *fewer* threads than
+    /// the budget, never more.
     #[arg(
         short = 'T',
         long,
@@ -83,9 +85,15 @@ pub struct QueryArgs {
 }
 
 impl QueryArgs {
+    /// `--threads` clamped to the process-wide CPU budget
+    /// (`obisys::cpu_budget`, bounded by the global `--cpu-max`).
+    pub fn effective_threads(&self) -> usize {
+        self.threads.min(obisys::cpu_budget()).max(1)
+    }
+
     pub fn effective_max_open(&self) -> usize {
         self.max_open_files
-            .unwrap_or_else(|| (self.threads / 4).max(1))
+            .unwrap_or_else(|| (self.effective_threads() / 4).max(1))
             .max(1)
     }
 }
@@ -131,7 +139,7 @@ pub fn run(args: QueryArgs) {
     let genomes = Arc::new(genomes);
     let n_partitions = idx.n_partitions();
     let with_counts = idx.meta().config.with_counts;
-    let n_workers = args.threads.max(1);
+    let n_workers = args.effective_threads();
 
     // Every partition/layer the query might touch is opened once, up front,
     // and shared (via Arc) across every `obipipeline` worker — a query pass

src/obikmer/src/cmd/superkmer/mod.rs

@@ -43,7 +43,7 @@ pub fn run(args: SuperkmerArgs) {
     let theta = args.common.theta;
     let level_max = args.common.level_max;
     let partition_bits = partitions_to_bits(args.common.partitions);
-    let n_workers = args.common.threads.max(1);
+    let n_workers = args.common.effective_threads();
     let max_open = args.common.effective_max_open();
 
     set_k(k);

src/obikmer/src/main.rs

@@ -2,11 +2,19 @@ mod cli;
 mod cmd;
 
 use clap::{Parser, Subcommand};
+use tracing::warn;
 use tracing_subscriber::{EnvFilter, fmt};
 
 #[derive(Parser)]
 #[command(name = "obikmer2", about = "DNA k-mer tools", version)]
 struct Cli {
+    /// Hard ceiling on the number of CPU cores the process may use — bounds
+    /// both the command's worker pool and every internal rayon pool.
+    /// Can only lower the budget, never raise it above the cores actually
+    /// available to the process. Defaults to that available count.
+    #[arg(long, global = true, value_name = "N")]
+    cpu_max: Option<usize>,
+
     #[command(subcommand)]
     command: Commands,
 }
@@ -53,6 +61,22 @@ fn main() {
         .init();
 
     let cli = Cli::parse();
+
+    // Install the CPU ceiling before anything sizes a thread pool.
+    if let Some(n) = cli.cpu_max {
+        obisys::set_cpu_cap(n);
+    }
+    if let Err(e) = rayon::ThreadPoolBuilder::new()
+        .num_threads(obisys::cpu_budget())
+        .build_global()
+    {
+        warn!("could not configure the global rayon pool: {e}");
+    }
+
+    // Always-on, near-zero-cost when disabled: RUST_LOG=obisys::profiler=trace
+    // to see periodic CPU/IO/memory snapshots on a production run.
+    obisys::start_profiler(std::time::Duration::from_secs(5));
+
     match cli.command {
         Commands::Index(args) => cmd::index::run(args),
         Commands::Superkmer(args) => cmd::superkmer::run(args),

src/obisys/src/lib.rs

@@ -4,6 +4,7 @@
 mod budget;
 mod lock;
 mod numa;
+mod profiler;
 mod progress;
 mod resources;
 mod stage;
@@ -11,6 +12,10 @@ mod stage;
 pub use budget::MemoryBudget;
 pub use lock::DirLock;
 pub use numa::PartitionRunner;
+pub use profiler::start as start_profiler;
 pub use progress::{Progress, TracedBar, progress_bar, spinner};
-pub use resources::{CpuSample, IoSample, available_memory_bytes, effective_parallelism, peak_rss_bytes};
+pub use resources::{
+    CpuSample, IoSample, available_memory_bytes, cpu_budget, effective_parallelism, peak_rss_bytes,
+    set_cpu_cap,
+};
 pub use stage::{Reporter, Stage, StageStats};

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
+/// each node's own worker cap, applied per node, so the pace of ramp-up
-/// ramp-up depends on node size rather than node count — a single-NUMA-node
+/// depends on that node's size rather than the node count — a
-/// (UMA) machine ramps just as fast as an 8-node one.
+/// 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 {
+            .map(|(pool, cpu_ids)| {
-                pool,
+                let max_workers = cpu_ids.len();
-                cpu_ids,
+                NodeConfig { pool, cpu_ids, max_workers }
-                max_workers: wpn,
             })
             .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
+    /// `(node's max_workers / INITIAL_DIVISOR).max(1)` are woken immediately
-    /// every node, then `(workers_per_node / GROWTH_DIVISOR).max(1)` more per
+    /// on every node, then `(node's max_workers / GROWTH_DIVISOR).max(1)`
-    /// node each time the check below fires. A timer thread fires that check
+    /// 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.
+    // UMA fallback: single synthetic node, budget-capped cores, no pool, no pinning.
-    let n_cores = crate::effective_parallelism();
+    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"))]
+/// Trims each NUMA node's CPU list, in place, so the total across all nodes
-pub fn build() -> NumaSetup {
+/// never exceeds `budget` — floor-split evenly across nodes. No-op when the
-    let n_cores = crate::effective_parallelism();
+/// topology already fits within `budget`. When `budget` is smaller than the
-    debug!("UMA: single synthetic node, {} core(s)", n_cores);
+/// number of nodes, the trailing nodes are emptied entirely (one core each
-    NumaSetup {
+/// to as many leading nodes as `budget` allows) rather than every node
-        pools: vec![None],
+/// keeping a token core that would collectively blow the budget.
-        cpus_per_node: vec![(0..n_cores).collect()],
+#[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();
+        }
     }
 }
 

src/obisys/src/profiler.rs

@@ -0,0 +1,50 @@
+use std::thread;
+use std::time::Duration;
+
+use tracing::{Level, trace};
+
+use crate::{CpuSample, IoSample, available_memory_bytes, cpu_budget, peak_rss_bytes};
+
+/// Spawn a detached background thread that logs a periodic snapshot of
+/// process-wide resource usage — CPU cores used (and as a % of the process's
+/// `--cpu-max` budget), I/O throughput, peak RSS, and remaining host/cgroup
+/// memory — at `trace!` level, every `interval`.
+///
+/// Meant to run for the whole process lifetime as always-on instrumentation
+/// installed once at startup (no stop handle). The sampling itself (a
+/// `getrusage` call and a `/proc/self/io` read) is cheap, but
+/// `available_memory_bytes()` builds a fresh `sysinfo::System` snapshot each
+/// time, so the whole body is skipped unless the `obisys::profiler` target
+/// is actually enabled at `TRACE`.
+///
+/// Enable with `RUST_LOG=obisys::profiler=trace` (combine with
+/// `RUST_LOG=debug,obisys::profiler=trace` to also get the per-partition
+/// `debug!` logs that already exist elsewhere) to profile a run directly on
+/// production hardware without attaching a separate profiler.
+pub fn start(interval: Duration) {
+    let budget = cpu_budget().max(1) as f64;
+    thread::Builder::new()
+        .name("obisys-profiler".into())
+        .spawn(move || {
+            let mut cpu = CpuSample::now();
+            let mut io = IoSample::now();
+            loop {
+                thread::sleep(interval);
+                if !tracing::enabled!(target: "obisys::profiler", Level::TRACE) {
+                    continue;
+                }
+                let cores = cpu.sample_cores().unwrap_or(0.0);
+                let io_mb_s = io.sample_rate().unwrap_or(0.0) / 1_000_000.0;
+                trace!(
+                    target: "obisys::profiler",
+                    cpu_pct = format_args!("{:.0}%", cores / budget * 100.0),
+                    cpu_cores = format_args!("{cores:.1}"),
+                    io_mb_s = format_args!("{io_mb_s:.1}"),
+                    rss_peak_mb = peak_rss_bytes() / 1_000_000,
+                    mem_avail_mb = available_memory_bytes() / 1_000_000,
+                    "profiler snapshot",
+                );
+            }
+        })
+        .expect("failed to spawn profiler thread");
+}

src/obisys/src/resources.rs

@@ -1,3 +1,4 @@
+use std::sync::OnceLock;
 use std::time::Instant;
 
 use libc::{RUSAGE_SELF, getrusage, rusage, timeval};
@@ -120,6 +121,32 @@ pub fn effective_parallelism() -> usize {
     host
 }
 
+// ── Process-wide CPU budget (hard ceiling) ───────────────────────────────────
+
+static CPU_CAP: OnceLock<usize> = OnceLock::new();
+
+/// Install a hard ceiling on the number of cores the process may use.
+///
+/// The value is clamped to `[1, effective_parallelism()]` — a cap can only
+/// lower the budget, never raise it above what the process is actually
+/// allowed to run on. First call wins; later calls are ignored.
+///
+/// Call this once at startup (before sizing any worker pool or the global
+/// rayon pool) when the user passes an explicit `--cpu-max`.
+pub fn set_cpu_cap(n: usize) {
+    let capped = n.clamp(1, effective_parallelism());
+    let _ = CPU_CAP.set(capped);
+}
+
+/// The process-wide CPU budget: the ceiling set via [`set_cpu_cap`] if any,
+/// otherwise [`effective_parallelism`].
+///
+/// Every worker pool and the global rayon pool must size themselves against
+/// this value rather than calling [`effective_parallelism`] directly.
+pub fn cpu_budget() -> usize {
+    CPU_CAP.get().copied().unwrap_or_else(effective_parallelism)
+}
+
 /// cgroup v2 (unified hierarchy): reads `cpu.max` ("<quota> <period>", or
 /// "max <period>" when unlimited) for the current process's cgroup, rounded
 /// up to whole cores. Returns `None` if unlimited or on any parse error.
@@ -237,18 +264,31 @@ impl CpuSample {
         cpu / (wall * n_cores as f64)
     }
 
-    pub fn do_i_activate(&mut self, threshold: f64) -> bool {
+    /// Cores of CPU time consumed per wall-clock second since the last call
+    /// (or since [`now`](Self::now) for the first call) — an absolute core
+    /// count, unlike [`cpu_efficiency`](Self::cpu_efficiency) which
+    /// normalizes to `n_cores` and never resets its baseline. Resets the
+    /// window on every call that returns `Some`; returns `None` (state left
+    /// untouched, so the window keeps accumulating) when less than 100 ms
+    /// have elapsed — too short to be meaningful.
+    pub fn sample_cores(&mut self) -> Option<f64> {
         let delta_wall = self.wall.elapsed().as_secs_f64();
         if delta_wall < 0.1 {
-            // Window too short to be meaningful — leave state untouched so it
+            return None;
-            // keeps accumulating until a real sample can be taken.
-            return false;
         }
-
         let n = CpuSample::now();
         let delta_ru = (n.user_secs - self.user_secs) + (n.sys_secs - self.sys_secs);
+        self.user_secs = n.user_secs;
+        self.sys_secs = n.sys_secs;
+        self.wall = n.wall;
+        Some(delta_ru / delta_wall)
+    }
+
+    pub fn do_i_activate(&mut self, threshold: f64) -> bool {
+        let Some(efficiency) = self.sample_cores() else {
+            return false;
+        };
 
-        let efficiency = delta_ru / delta_wall;
         let activate = 0f64.max(efficiency - self.previous) >= threshold;
 
         debug!(
@@ -259,9 +299,6 @@ impl CpuSample {
             activate
         );
         self.previous = efficiency;
-        self.user_secs = n.user_secs;
-        self.sys_secs = n.sys_secs;
-        self.wall = n.wall;
 
         activate
     }
@@ -324,18 +361,32 @@ impl IoSample {
         0
     }
 
+    /// Bytes/second submitted to the block layer since the last call (or
+    /// since [`now`](Self::now) for the first call). Resets the window on
+    /// every call that returns `Some`; returns `None` (state left
+    /// untouched) when less than 100 ms have elapsed — too short to be
+    /// meaningful.
+    pub fn sample_rate(&mut self) -> Option<f64> {
+        let elapsed = self.wall.elapsed().as_secs_f64();
+        if elapsed < 0.1 {
+            return None;
+        }
+        let n = Self::read_bytes();
+        let rate = n.saturating_sub(self.bytes) as f64 / elapsed;
+        self.bytes = n;
+        self.wall = Instant::now();
+        Some(rate)
+    }
+
     /// Same protocol as [`CpuSample::do_i_activate`] (0.1 s minimum window,
     /// state untouched on early return), but growth is measured relative to
     /// the previous rate. `threshold` is a fraction, e.g. `0.2` for a 20 %
     /// increase in throughput since the last real sample.
     pub fn do_i_activate(&mut self, threshold: f64) -> bool {
-        let elapsed = self.wall.elapsed().as_secs_f64();
+        let Some(rate) = self.sample_rate() else {
-        if elapsed < 0.1 {
             return false;
-        }
+        };
 
-        let n = Self::read_bytes();
-        let rate = n.saturating_sub(self.bytes) as f64 / elapsed;
         let activate = if self.previous_rate == 0.0 {
             rate > 0.0 // bootstrap: any measured throughput is signal enough
         } else {
@@ -347,8 +398,6 @@ impl IoSample {
             self.previous_rate, rate, activate
         );
         self.previous_rate = rate;
-        self.bytes = n;
-        self.wall = Instant::now();
 
         activate
     }