diff --git a/docs/roadmap/svar-2.md b/docs/roadmap/svar-2.md index 4d5e449..42622fd 100644 --- a/docs/roadmap/svar-2.md +++ b/docs/roadmap/svar-2.md @@ -401,9 +401,11 @@ Legend: `[ ]` not started · `[~]` in progress · `[x]` done so Approach B was not built. Public API is unchanged — sharding is entirely internal to the existing `threads=` budget. Byte-identical output vs. serial conversion is gated by a store-hash oracle at every thread count for both - backends. **VCF** scales ~3.9× at 32 cores (chr21, 1176s → 300s); sub-contig - sharding only engages once the thread budget clears HTSlib decode-thread - allocation (~15 cores). **PGEN** sharding is byte-identical but not faster — + backends. **VCF** scales ~3.9× at 32 cores (chr21, 1176s → 300s). Its + backend-specific `reader_workers` budget treats indexed shard readers as a + replacement for the monolithic reader's HTSlib pool and disables per-shard + HTSlib background pools, so medium-sized single-contig runs no longer strand + cores in an inactive reservation. **PGEN** sharding is byte-identical but not faster — `pgenlib`'s genotype decode holds the CPython GIL, so shard readers serialize and sharding is net slightly slower than serial (memory `pgenlib-holds-gil-sharded-reads`); PGEN is intentionally not over-decomposed. diff --git a/docs/source/svar.md b/docs/source/svar.md index 3dd34bf..e44d109 100644 --- a/docs/source/svar.md +++ b/docs/source/svar.md @@ -73,13 +73,14 @@ regions per contig raise — use `pos`/`record`, or convert separately. Single-file `SparseVar2.from_vcf` shards **within a contig**, driven by the same `threads=` budget shown above — no new argument. Sub-contig sharding only kicks -in for the default whole-contig (`regions_overlap="pos"`) path: the thread budget -first spends added cores on HTSlib decode threads for the single reader, so the -sub-contig shard budget stays at 1 (an un-sharded reader) until the core count -clears that stage (~15 cores on the benchmarked hardware). Output is -**byte-identical** to serial conversion at every thread count — sharding is gated -by a store-hash oracle and does not reintroduce missingness (a `./.` haplotype -and a hom-ref haplotype remain indistinguishable in SVAR2 either way). +in for the default whole-contig (`regions_overlap="pos"`) path. The planner uses +a backend-specific reader budget: indexed shard readers decompress inline and +replace, rather than run alongside, the monolithic reader's HTSlib pool. This +lets medium-sized single-contig runs use their available cores without +oversubscribing multi-contig runs. Output is **byte-identical** to serial +conversion at every thread count — sharding is gated by a store-hash oracle and +does not reintroduce missingness (a `./.` haplotype and a hom-ref haplotype +remain indistinguishable in SVAR2 either way). Sub-contig sharding is restricted to `regions_overlap="pos"` (which the whole-contig default uses). `"record"` and `"variant"` conversions run on a diff --git a/src/budget.rs b/src/budget.rs index a31ece0..797e1b1 100644 --- a/src/budget.rs +++ b/src/budget.rs @@ -4,6 +4,10 @@ // 4 fixed OS threads per chrom: reader + executor + chunk_writer + long_allele_writer. pub const PIPELINE_THREADS_PER_CHROM: usize = 4; +// Independent indexed VCF shard readers decompress in their worker thread. +// Giving each one an HTSlib background pool would multiply the process-wide +// thread budget by the shard count. +pub const SHARDED_VCF_HTSLIB_THREADS_PER_READER: usize = 0; // Floor for HTSlib decode threads — below this the executor channel starves. const MIN_HTSLIB_THREADS: usize = 2; // Ceiling for HTSlib decode threads. Bumped 4→8 for single-/few-contig @@ -18,10 +22,13 @@ const MIN_THREADS_PER_CHROM: usize = PIPELINE_THREADS_PER_CHROM + MIN_HTSLIB_THR pub struct ThreadPlan { pub concurrent_chroms: usize, pub htslib_threads: usize, + // Indexed VCF shard readers per concurrent contig. Unlike + // `processing_threads`, this budget reclaims the monolithic reader's + // HTSlib pool because each shard decompresses inline. + pub reader_workers: usize, // Cores left idle after the pipeline + htslib threads across all concurrent - // chroms. For splittable VCF contigs this caps concurrent shard readers; - // otherwise it sizes the reader-side processing pool used for bounded - // normalization batches plus intra-chunk presence packing. + // chroms. This sizes the non-sharded reader-side processing pool used for + // bounded normalization batches plus intra-chunk presence packing. pub processing_threads: usize, } @@ -40,6 +47,7 @@ pub fn plan_thread_budget(available_cores: usize, n_chroms: usize) -> ThreadPlan ThreadPlan { concurrent_chroms: 1, htslib_threads: htslib, + reader_workers: reader_workers(usable_cores, 1), processing_threads: processing, } } else { @@ -53,6 +61,7 @@ pub fn plan_thread_budget(available_cores: usize, n_chroms: usize) -> ThreadPlan ThreadPlan { concurrent_chroms: concurrent, htslib_threads: htslib, + reader_workers: reader_workers(usable_cores, concurrent), processing_threads: processing, } } @@ -65,6 +74,21 @@ fn processing_threads(usable_cores: usize, concurrent: usize, htslib: usize) -> usable_cores.saturating_sub(active).max(1) } +/// Per-contig worker count for the indexed/sharded VCF backend. +/// +/// Shard readers replace the monolithic reader's HTSlib pool rather than +/// running alongside it, so split the usable process budget evenly across +/// active contigs and spend the remainder after their fixed pipeline threads. +fn reader_workers(usable_cores: usize, concurrent: usize) -> usize { + let cores_per_chrom = usable_cores / concurrent.max(1); + let worker_cost = 1 + SHARDED_VCF_HTSLIB_THREADS_PER_READER; + cores_per_chrom + .saturating_sub(PIPELINE_THREADS_PER_CHROM) + .checked_div(worker_cost) + .unwrap_or(0) + .max(1) +} + #[cfg(test)] mod tests { use super::*; @@ -76,6 +100,7 @@ mod tests { ThreadPlan { concurrent_chroms: 1, htslib_threads: 1, + reader_workers: 1, processing_threads: 1, } ); @@ -88,6 +113,7 @@ mod tests { ThreadPlan { concurrent_chroms: 1, htslib_threads: 1, + reader_workers: 1, processing_threads: 1, } ); @@ -100,6 +126,7 @@ mod tests { ThreadPlan { concurrent_chroms: 10, htslib_threads: 2, + reader_workers: 2, processing_threads: 4, } ); @@ -142,6 +169,34 @@ mod tests { assert_eq!(plan.processing_threads, 20); } + #[test] + fn test_sharded_vcf_reclaims_unused_htslib_budget_for_reader_workers() { + // Sharded VCF readers each use one inline HTSlib thread, so the separate + // 8-thread HTSlib decode pool is not active on this backend. A 16-core, + // one-contig run therefore has 15 usable cores: 4 fixed pipeline threads + // plus 11 independent shard readers. + let plan = plan_thread_budget(16, 1); + assert_eq!(plan.reader_workers, 11); + assert_eq!( + plan.concurrent_chroms + * (PIPELINE_THREADS_PER_CHROM + + plan.reader_workers * (1 + SHARDED_VCF_HTSLIB_THREADS_PER_READER)), + 15 + ); + } + + #[test] + fn test_sharded_vcf_reader_workers_are_bounded_across_concurrent_contigs() { + // 65 cores → 64 usable; 10 concurrent contigs × (4 fixed + 2 readers) + // = 60 active sharded-path threads, leaving four cores of headroom. + let plan = plan_thread_budget(65, 22); + let active = plan.concurrent_chroms + * (PIPELINE_THREADS_PER_CHROM + + plan.reader_workers * (1 + SHARDED_VCF_HTSLIB_THREADS_PER_READER)); + assert_eq!(plan.reader_workers, 2); + assert!(active <= 64); + } + #[test] fn test_processing_threads_floored_at_one_when_saturated() { // 65 cores → usable 64; 22 chroms → concurrent 10; htslib 2. diff --git a/src/lib.rs b/src/lib.rs index cf7e978..fc1eb44 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -222,15 +222,21 @@ fn run_conversion_pipeline( let plan = crate::budget::plan_thread_budget(available_cores, chroms.len()); let concurrent_chroms = plan.concurrent_chroms; let htslib_threads = plan.htslib_threads; + let reader_workers = plan.reader_workers; let processing_threads = plan.processing_threads; - let total_active = + let monolithic_reader_active = concurrent_chroms * (crate::budget::PIPELINE_THREADS_PER_CHROM + htslib_threads); + let sharded_vcf_active = concurrent_chroms + * (crate::budget::PIPELINE_THREADS_PER_CHROM + + reader_workers * (1 + crate::budget::SHARDED_VCF_HTSLIB_THREADS_PER_READER)); tracing::info!(cores = available_cores, "using cores"); tracing::info!( concurrent_chroms, htslib_threads, - total_active, + monolithic_reader_active, + reader_workers, + sharded_vcf_active, processing_threads, "pipeline config" ); @@ -257,6 +263,7 @@ fn run_conversion_pipeline( orchestrator::SourceSpec::Vcf { vcf_path: vcf_path.clone(), htslib_threads, + reader_workers, regions: ranges_by_chrom.get(chrom).cloned().unwrap_or_default(), overlap: overlap_mode, }, diff --git a/src/orchestrator.rs b/src/orchestrator.rs index d98eb15..97ba5a3 100644 --- a/src/orchestrator.rs +++ b/src/orchestrator.rs @@ -63,6 +63,9 @@ pub enum SourceSpec { Vcf { vcf_path: String, htslib_threads: usize, + /// Independent indexed shard readers for this contig. These replace + /// the monolithic reader's HTSlib pool on the sharded path. + reader_workers: usize, regions: Vec<(u32, u32)>, overlap: crate::svar2_view::OverlapMode, }, @@ -384,6 +387,7 @@ pub fn process_chromosome( SourceSpec::Vcf { vcf_path, htslib_threads, + reader_workers, regions, overlap, } => { @@ -409,7 +413,7 @@ pub fn process_chromosome( crate::vcf_reader::plan_vcf_shards( ®ions, &chr, - processing_threads.saturating_mul(OVERSHARD_FACTOR), + reader_workers.saturating_mul(OVERSHARD_FACTOR), chunk_size as u32, )? } else { @@ -435,19 +439,19 @@ pub fn process_chromosome( .collect(); trace_ll!( "[plan {chr}] workers={} shards={}", - processing_threads, + reader_workers, units.len() ); let totals = crate::shard_exec::run( &chr, units, - processing_threads, + reader_workers, |unit| { let source = crate::vcf_reader::VcfRecordSource::new( &vcf_path, &chr, &s_refs, - 1, // htslib_threads: many concurrent shard readers, keep each small + crate::budget::SHARDED_VCF_HTSLIB_THREADS_PER_READER, ploidy, &fields_owned, // The shard's padded fetch window IS the @@ -748,8 +752,9 @@ pub fn process_chromosome( // Dedicated rayon pool for reader-side CPU work: bounded per-record // normalization batches plus intra-chunk presence packing. The // sharded VCF branch above returns before this point because its - // independent indexed readers consume the same `processing_threads` - // budget directly; building both would double-reserve cores. + // independent indexed readers consume the backend-specific + // `reader_workers` budget directly; building both would + // double-reserve cores. let pool = rayon::ThreadPoolBuilder::new() .num_threads(processing_threads.max(1)) .thread_name(|i| format!("pack-{}", i)) diff --git a/tests/common/mod.rs b/tests/common/mod.rs index 3794503..a213606 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -177,6 +177,7 @@ pub fn build_contig( genoray_core::orchestrator::SourceSpec::Vcf { vcf_path: bcf.to_str().unwrap().to_string(), htslib_threads: 1, + reader_workers: 1, regions: Vec::new(), overlap: genoray_core::svar2_view::OverlapMode::Pos, }, diff --git a/tests/test_check_ref_e2e.rs b/tests/test_check_ref_e2e.rs index 9b8d66e..70ed62b 100644 --- a/tests/test_check_ref_e2e.rs +++ b/tests/test_check_ref_e2e.rs @@ -70,6 +70,7 @@ fn convert( SourceSpec::Vcf { vcf_path: bcf.to_str().unwrap().to_string(), htslib_threads: 1, + reader_workers: 1, regions: Vec::new(), overlap: genoray_core::svar2_view::OverlapMode::Pos, }, @@ -197,7 +198,7 @@ fn vcf_list_ref_mismatch_excluded_under_x() { // region list disables sharding — see the Python `from_vcf` comment on why // it always fills `[0, len)` for whole-contig conversion) and // `overlap == OverlapMode::Pos`. Passing the whole-contig range explicitly, -// a small `chunk_size` (target shard span), and `processing_threads > 1` +// a small `chunk_size` (target shard span), and `reader_workers > 1` // reproduces the Python test's sharded scenario (there: `threads=16, // chunk_size=1`) directly against the Rust entry point. #[test] @@ -218,6 +219,7 @@ fn sharded_ref_excluded_counted_once_in_contig_done() { SourceSpec::Vcf { vcf_path: bcf.to_str().unwrap().to_string(), htslib_threads: 1, + reader_workers: 8, // Non-empty, whole-contig range: required to enable sub-contig // sharding (see comment above). regions: vec![(0, 1000)], diff --git a/tests/test_convert_skip_e2e.rs b/tests/test_convert_skip_e2e.rs index a3765f9..bb4fbda 100644 --- a/tests/test_convert_skip_e2e.rs +++ b/tests/test_convert_skip_e2e.rs @@ -44,6 +44,7 @@ fn convert( genoray_core::orchestrator::SourceSpec::Vcf { vcf_path: bcf.to_str().unwrap().to_string(), htslib_threads: 1, + reader_workers: 1, regions: Vec::new(), overlap: genoray_core::svar2_view::OverlapMode::Pos, }, diff --git a/tests/test_e2e.rs b/tests/test_e2e.rs index e01dbd9..d993fc9 100644 --- a/tests/test_e2e.rs +++ b/tests/test_e2e.rs @@ -93,6 +93,7 @@ fn test_e2e_normalized_bcf_pipeline() { genoray_core::orchestrator::SourceSpec::Vcf { vcf_path: bcf_path.to_str().unwrap().to_string(), htslib_threads: 1, + reader_workers: 1, regions: Vec::new(), overlap: genoray_core::svar2_view::OverlapMode::Pos, }, @@ -265,6 +266,7 @@ fn test_e2e_max_del_postpass() { genoray_core::orchestrator::SourceSpec::Vcf { vcf_path: bcf_path.to_str().unwrap().to_string(), htslib_threads: 1, + reader_workers: 1, regions: Vec::new(), overlap: genoray_core::svar2_view::OverlapMode::Pos, }, @@ -346,6 +348,7 @@ fn test_e2e_dense_snp_roundtrip() { genoray_core::orchestrator::SourceSpec::Vcf { vcf_path: bcf_path.to_str().unwrap().to_string(), htslib_threads: 1, + reader_workers: 1, regions: Vec::new(), overlap: genoray_core::svar2_view::OverlapMode::Pos, }, @@ -426,6 +429,7 @@ fn test_e2e_mutation_conservation() { genoray_core::orchestrator::SourceSpec::Vcf { vcf_path: bcf_path.to_str().unwrap().to_string(), htslib_threads: 1, + reader_workers: 1, regions: Vec::new(), overlap: genoray_core::svar2_view::OverlapMode::Pos, }, @@ -886,6 +890,7 @@ fn test_missing_chrom_returns_err() { genoray_core::orchestrator::SourceSpec::Vcf { vcf_path: bcf_path.to_str().unwrap().to_string(), htslib_threads: 1, + reader_workers: 1, regions: Vec::new(), overlap: genoray_core::svar2_view::OverlapMode::Pos, }, @@ -996,6 +1001,7 @@ fn regions_overlap_variant_keeps_spanning_deletion_e2e() { genoray_core::orchestrator::SourceSpec::Vcf { vcf_path: bcf_path.to_str().unwrap().to_string(), htslib_threads: 1, + reader_workers: 1, regions: vec![(6, 12)], overlap, },