From b66a488277bbf87e512afe80d1ab7bd805799a46 Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Thu, 10 Sep 2026 17:07:14 +0200 Subject: [PATCH] Introduce CPU budget and dynamic thread limiting This change implements mechanisms to enforce resource limits based on the available CPU budget, including setting a hard CPU ceiling and calculating effective thread counts for CLI, query, and superkmer operations. --- src/Cargo.lock | 3 ++- src/obikmer/Cargo.toml | 3 ++- src/obikmer/src/cli.rs | 13 +++++++++++-- src/obikmer/src/cmd/index/mod.rs | 2 +- src/obikmer/src/cmd/query/mod.rs | 14 +++++++++++--- src/obikmer/src/cmd/superkmer/mod.rs | 2 +- src/obikmer/src/main.rs | 20 ++++++++++++++++++++ src/obisys/src/lib.rs | 5 ++++- src/obisys/src/resources.rs | 27 +++++++++++++++++++++++++++ 9 files changed, 79 insertions(+), 10 deletions(-) diff --git a/src/Cargo.lock b/src/Cargo.lock index a3c8de48..20ddaff7 100644 --- a/src/Cargo.lock +++ b/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", diff --git a/src/obikmer/Cargo.toml b/src/obikmer/Cargo.toml index f7e2e9e4..84c8fcb5 100644 --- a/src/obikmer/Cargo.toml +++ b/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"] } diff --git a/src/obikmer/src/cli.rs b/src/obikmer/src/cli.rs index 6018ad30..8a1abb06 100644 --- a/src/obikmer/src/cli.rs +++ b/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) } diff --git a/src/obikmer/src/cmd/index/mod.rs b/src/obikmer/src/cmd/index/mod.rs index 7bcace91..6d178323 100644 --- a/src/obikmer/src/cmd/index/mod.rs +++ b/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"); diff --git a/src/obikmer/src/cmd/query/mod.rs b/src/obikmer/src/cmd/query/mod.rs index a86a6583..e9526b5b 100644 --- a/src/obikmer/src/cmd/query/mod.rs +++ b/src/obikmer/src/cmd/query/mod.rs @@ -63,7 +63,9 @@ pub struct QueryArgs { #[arg(short = 'z', long)] pub findere_z: Option, - /// 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 diff --git a/src/obikmer/src/cmd/superkmer/mod.rs b/src/obikmer/src/cmd/superkmer/mod.rs index 707651f9..a6d9edf1 100644 --- a/src/obikmer/src/cmd/superkmer/mod.rs +++ b/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); diff --git a/src/obikmer/src/main.rs b/src/obikmer/src/main.rs index d56da73f..ae68c132 100644 --- a/src/obikmer/src/main.rs +++ b/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, + #[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), diff --git a/src/obisys/src/lib.rs b/src/obisys/src/lib.rs index a37cceff..63da76f8 100644 --- a/src/obisys/src/lib.rs +++ b/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}; diff --git a/src/obisys/src/resources.rs b/src/obisys/src/resources.rs index 4ea51f09..6f8d4ad0 100644 --- a/src/obisys/src/resources.rs +++ b/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 = 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` (" ", or /// "max " when unlimited) for the current process's cgroup, rounded /// up to whole cores. Returns `None` if unlimited or on any parse error. -- 2.54.0