feat(query): add throughput metering and --max-open-files flag

Introduces an EMA-based throughput meter that dynamically updates a spinner with MB/s rates, along with atomic counters for tracking cumulative bytes and active chunks. Adds final pipeline reporting and consolidates imports for cleaner performance instrumentation.
This commit is contained in:
Eric Coissac
2026-07-07 15:19:09 +02:00
parent 8f0ceec784
commit 61c390503d
+77 -4
View File
@@ -2,6 +2,8 @@ use std::collections::{HashMap, VecDeque};
use std::io::{self, BufWriter, Write};
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
use std::time::Instant;
use clap::Args;
use obikindex::KmerIndex;
@@ -11,8 +13,8 @@ use obilayeredmap::IndexMode;
use obiread::chunk::read_sequence_chunks_sized;
use obiread::record::{SeqRecord, parse_chunk};
use obiskbuilder::SuperKmerIter;
use obisys::available_memory_bytes;
use tracing::info;
use obisys::{Reporter, Stage, available_memory_bytes, spinner};
use tracing::{debug, info};
// ── Pipeline data ─────────────────────────────────────────────────────────────
@@ -74,6 +76,20 @@ pub struct QueryArgs {
/// I/O chunk size in MiB (default: auto-sized from available RAM and thread count)
#[arg(long)]
pub chunk_size: Option<usize>,
/// Maximum number of input files open simultaneously.
/// Defaults to threads/4 (minimum 1). Keep below the number of workers
/// to ensure CPU workers are always available for the transform stage.
#[arg(long)]
pub max_open_files: Option<usize>,
}
impl QueryArgs {
pub fn effective_max_open(&self) -> usize {
self.max_open_files
.unwrap_or_else(|| (self.threads / 4).max(1))
.max(1)
}
}
// ── SKDesc — one occurrence of a superkmer in the batch ───────────────────────
@@ -234,6 +250,9 @@ fn process_chunk(
force_presence: bool,
presence_threshold: u32,
) -> Vec<u8> {
let chunk_start = Instant::now();
let chunk_bytes = rope.len();
let records = parse_chunk(&rope, k);
if records.is_empty() {
return Vec::new();
@@ -384,6 +403,15 @@ fn process_chunk(
&cov,
&mut buf,
);
debug!(
chunk_bytes,
n_seqs,
n_smers = batch.n_kmers.iter().map(|&n| n as u64).sum::<u64>(),
wall_ms = chunk_start.elapsed().as_millis() as u64,
"process_chunk"
);
buf
}
@@ -457,26 +485,71 @@ pub fn run(args: QueryArgs) {
}
});
// Instrumentation: total bytes processed (for the EMA throughput readout)
// and the number of chunks currently being processed by a worker — both
// read from the spinner loop below, updated from inside the pipe closure.
let total_bytes = Arc::new(AtomicU64::new(0));
let chunks_active = Arc::new(AtomicU32::new(0));
let pipe = obipipeline::make_pipe! {
QueryData : Rope => Vec<u8>,
| {
let idx = Arc::clone(&idx);
let total_bytes = Arc::clone(&total_bytes);
let chunks_active = Arc::clone(&chunks_active);
move |rope: Rope| {
process_chunk(
chunks_active.fetch_add(1, Ordering::Relaxed);
let bytes = rope.len() as u64;
let out = process_chunk(
&idx, rope, k, n_genomes, n_partitions, with_counts,
effective_z, detail, count_missing, force_presence, presence_threshold,
)
);
total_bytes.fetch_add(bytes, Ordering::Relaxed);
chunks_active.fetch_sub(1, Ordering::Relaxed);
out
}
} : Chunk => Output,
};
let t = Stage::start("query");
let pb = spinner("query");
let mut ema_rate: f64 = 0.0;
let mut last_t = Instant::now();
let mut last_bytes: u64 = 0;
const ALPHA: f64 = 0.15;
let mut out = BufWriter::new(io::stdout());
for block in pipe.apply(all_chunks, n_workers, 2) {
if !block.is_empty() {
out.write_all(&block).expect("write error");
}
let now = Instant::now();
let dt = now.duration_since(last_t).as_secs_f64();
if dt > 0.1 {
let total = total_bytes.load(Ordering::Relaxed);
let instant = (total - last_bytes) as f64 / dt;
ema_rate = ALPHA * instant + (1.0 - ALPHA) * ema_rate;
last_t = now;
last_bytes = total;
let bp = total as f64;
let (count_str, rate_str) = if bp >= 1e9 {
(format!("{:.2} GB", bp / 1e9), format!("{:.0} MB/s", ema_rate / 1e6))
} else {
(format!("{:.0} MB", bp / 1e6), format!("{:.0} MB/s", ema_rate / 1e6))
};
let active = chunks_active.load(Ordering::Relaxed);
pb.set_message(format!("{count_str} {rate_str} [chunks in flight: {active}]"));
}
}
out.flush().expect("flush error");
pb.finish_and_clear();
let mut rep = Reporter::new();
rep.push(t.stop());
rep.print();
}
// ── Output ────────────────────────────────────────────────────────────────────