Merge pull request 'Introduce CPU budget and dynamic thread limiting' (#73) from push-lztkokoxqpuw into main

Reviewed-on: #73

src/Cargo.lock

@@ -1480,7 +1480,7 @@ dependencies = [
 
 [[package]]
 name = "obikmer"
-version = "1.3.1"
+version = "1.3.2"
 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.2"
 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,18 @@ 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}");
+    }
+
     match cli.command {
         Commands::Index(args) => cmd::index::run(args),
         Commands::Superkmer(args) => cmd::superkmer::run(args),

src/obisys/src/lib.rs

@@ -12,5 +12,8 @@ pub use budget::MemoryBudget;
 pub use lock::DirLock;
 pub use numa::PartitionRunner;
 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/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.