diff --git a/src/obikmer/src/cmd/query.rs b/src/obikmer/src/cmd/query.rs index b0fe090..4972c13 100644 --- a/src/obikmer/src/cmd/query.rs +++ b/src/obikmer/src/cmd/query.rs @@ -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, + + /// 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, +} + +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 { + 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::(), + 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, | { 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 ────────────────────────────────────────────────────────────────────