Skip to main content

cuttlefish_rs/
params.rs

1//! Build configuration and phase-local resource planning.
2//!
3//! [`BuildParams::threads`] is an upper bound. When `max_memory_gb` is set,
4//! partition, local contraction, and post-local processing independently lower
5//! concurrency according to replicated state estimates. The limit is soft:
6//! workload-sized shared tables can impose a higher minimum RSS.
7
8use crate::{
9    DEFAULT_CUTOFF_READS, DEFAULT_CUTOFF_REFS, DEFAULT_GMTIG_BUCKETS, DEFAULT_K,
10    DEFAULT_LMTIG_BUCKETS, DEFAULT_MINIMIZER_LEN, DEFAULT_VERTEX_PARTITIONS, GraphInput, MAX_K,
11    MAX_MINIMIZER_LEN, default_threads, default_work_dir,
12};
13
14const GIB: usize = 1024 * 1024 * 1024;
15// This is the soft-memory policy for the Rust-only `--max-memory` extension,
16// not a C++ algorithm parameter. These estimates cover phase-local replicated
17// state; workload-sized graph tables can impose a higher minimum.
18const PARTITION_FIXED_MEMORY: usize = 2 * GIB;
19const PARTITION_MEMORY_PER_WORKER: usize = 32 * 1024 * 1024;
20const LOCAL_FIXED_MEMORY: usize = 4 * GIB;
21const LOCAL_MEMORY_PER_WORKER: usize = 160 * 1024 * 1024;
22const POST_LOCAL_FIXED_MEMORY: usize = 4 * GIB;
23const POST_LOCAL_MEMORY_PER_WORKER: usize = 64 * 1024 * 1024;
24
25/// Configuration shared by partitioning and graph construction phases.
26///
27/// Bucket and partition counts must be powers of two. The public fields support
28/// programmatic construction, while [`BuildParams::validate`] enforces the
29/// cross-field invariants expected by the production pipeline.
30#[derive(Debug, Clone, PartialEq, Eq)]
31pub struct BuildParams {
32    /// Whether inputs are sequencing reads or references.
33    pub input: GraphInput,
34    /// Explicit sequence file paths.
35    pub seqs: Vec<String>,
36    /// Files containing one input path per line.
37    pub lists: Vec<String>,
38    /// Directories whose regular files are used as inputs.
39    pub dirs: Vec<String>,
40    /// K-mer length. Must be odd and no greater than [`crate::MAX_K`].
41    pub k: u16,
42    /// Minimizer length used for weak-super-k-mer partitioning.
43    pub minimizer_len: u16,
44    /// Optional edge-frequency cutoff; input-class defaults apply when absent.
45    pub cutoff: Option<u32>,
46    /// Whether to construct and emit positional colors.
47    pub color: bool,
48    /// Whether weak-super-k-mer buckets use LZ4 block compression.
49    ///
50    /// On by default. Colored builds compress regardless; for uncolored ones
51    /// this cuts write volume on a 150k reference build from 1167.7 GB to
52    /// 439 GB and peak RSS by about a quarter, at no measurable cost in wall
53    /// time across five interleaved pairs at 64 threads.
54    pub compress_buckets: bool,
55    /// Whether an input that fails to parse is skipped instead of aborting.
56    ///
57    /// A skipped source retains its position in the input list, so colored
58    /// source assignments are unaffected. Records read before the failure are
59    /// retained, so a file that fails part-way through contributes its prefix.
60    pub skip_unreadable: bool,
61    /// Prefix for final graph files.
62    pub output_prefix: String,
63    /// Directory for external-memory intermediates.
64    pub work_dir: String,
65    /// Number of blocked discontinuity-graph vertex partitions.
66    pub vertex_partitions: usize,
67    /// Logical local-unitig bucket count.
68    pub lmtig_buckets: usize,
69    /// Logical maximal-unitig coordinate bucket count.
70    pub gmtig_buckets: usize,
71    /// Maximum worker count requested for each phase.
72    pub threads: usize,
73    /// Optional soft memory budget in GiB.
74    ///
75    /// The budget controls replicated phase state and is not a hard RSS limit.
76    pub max_memory_gb: Option<usize>,
77}
78
79impl BuildParams {
80    /// Constructs parameters with Cuttlefish-compatible defaults.
81    pub fn new(input: GraphInput, output_prefix: String) -> Self {
82        Self {
83            input,
84            seqs: Vec::new(),
85            lists: Vec::new(),
86            dirs: Vec::new(),
87            k: DEFAULT_K,
88            minimizer_len: DEFAULT_MINIMIZER_LEN,
89            cutoff: None,
90            color: false,
91            compress_buckets: true,
92            skip_unreadable: false,
93            output_prefix,
94            work_dir: default_work_dir(),
95            vertex_partitions: DEFAULT_VERTEX_PARTITIONS,
96            lmtig_buckets: DEFAULT_LMTIG_BUCKETS,
97            gmtig_buckets: DEFAULT_GMTIG_BUCKETS,
98            threads: default_threads(),
99            max_memory_gb: None,
100        }
101    }
102
103    /// Returns the explicit cutoff or the input-class default.
104    #[inline]
105    pub fn cutoff(&self) -> u32 {
106        self.cutoff.unwrap_or_else(|| self.input.default_cutoff())
107    }
108
109    /// Returns the soft memory budget in bytes, if configured.
110    pub fn max_memory_bytes(&self) -> Option<usize> {
111        self.max_memory_gb.and_then(|gb| gb.checked_mul(GIB))
112    }
113
114    fn memory_bounded_workers(&self, fixed: usize, per_worker: usize) -> usize {
115        let requested = self.threads.max(1);
116        let Some(budget) = self.max_memory_bytes() else {
117            return requested;
118        };
119        let bounded = requested
120            .min(budget.saturating_sub(fixed) / per_worker)
121            .max(1);
122        if bounded == requested {
123            return requested;
124        }
125        // When the estimate requires a reduction, round down to leave headroom
126        // for workload-sized shared tables outside this replicated-state estimate.
127        1usize << bounded.ilog2()
128    }
129
130    /// Returns the partition worker count after input and memory bounds.
131    pub fn partition_workers(&self, input_files: usize) -> usize {
132        self.memory_bounded_workers(PARTITION_FIXED_MEMORY, PARTITION_MEMORY_PER_WORKER)
133            .min(input_files.max(1))
134    }
135
136    /// Returns the memory-bounded local-contraction worker count.
137    pub fn local_workers(&self) -> usize {
138        self.memory_bounded_workers(LOCAL_FIXED_MEMORY, LOCAL_MEMORY_PER_WORKER)
139    }
140
141    /// Returns the memory-bounded discontinuity/collation worker count.
142    pub fn post_local_workers(&self) -> usize {
143        self.memory_bounded_workers(POST_LOCAL_FIXED_MEMORY, POST_LOCAL_MEMORY_PER_WORKER)
144    }
145
146    /// Validates input, graph dimensions, and resource parameters.
147    pub fn validate(&self) -> Result<(), ParamError> {
148        if self.seqs.is_empty() && self.lists.is_empty() && self.dirs.is_empty() {
149            return Err(ParamError::NoInput);
150        }
151
152        if self.output_prefix.is_empty() {
153            return Err(ParamError::MissingOutput);
154        }
155
156        if self.k <= 1 || self.k > MAX_K || self.k % 2 == 0 {
157            return Err(ParamError::InvalidK(self.k));
158        }
159
160        if self.minimizer_len == 0
161            || self.minimizer_len >= self.k
162            || self.minimizer_len > MAX_MINIMIZER_LEN
163        {
164            return Err(ParamError::InvalidMinimizerLen {
165                k: self.k,
166                l: self.minimizer_len,
167            });
168        }
169
170        if self.cutoff() == 0 {
171            return Err(ParamError::InvalidCutoff);
172        }
173
174        if !self.vertex_partitions.is_power_of_two()
175            || !self.lmtig_buckets.is_power_of_two()
176            || !self.gmtig_buckets.is_power_of_two()
177        {
178            return Err(ParamError::BucketCountsMustBePowersOfTwo);
179        }
180        if self.threads == 0 {
181            return Err(ParamError::InvalidThreadCount);
182        }
183        if self.max_memory_gb == Some(0)
184            || self.max_memory_gb.is_some_and(|gb| gb > usize::MAX / GIB)
185        {
186            return Err(ParamError::InvalidMemoryLimit);
187        }
188
189        Ok(())
190    }
191}
192
193#[derive(Debug, Clone, Copy, PartialEq, Eq)]
194pub enum ParamError {
195    NoInput,
196    MissingOutput,
197    InvalidK(u16),
198    InvalidMinimizerLen { k: u16, l: u16 },
199    InvalidCutoff,
200    BucketCountsMustBePowersOfTwo,
201    InvalidThreadCount,
202    InvalidMemoryLimit,
203}
204
205impl std::fmt::Display for ParamError {
206    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
207        match self {
208            Self::NoInput => write!(f, "no sequence input provided"),
209            Self::MissingOutput => write!(f, "missing output prefix"),
210            Self::InvalidK(k) => write!(f, "k-mer length {k} is invalid; expected odd 3..={MAX_K}"),
211            Self::InvalidMinimizerLen { k, l } => {
212                write!(f, "minimizer length {l} is invalid for k={k}")
213            }
214            Self::InvalidCutoff => write!(f, "cutoff frequency must be at least 1"),
215            Self::BucketCountsMustBePowersOfTwo => write!(f, "bucket counts must be powers of two"),
216            Self::InvalidThreadCount => write!(f, "thread count must be at least 1"),
217            Self::InvalidMemoryLimit => write!(f, "maximum memory must be at least 1 GiB"),
218        }
219    }
220}
221
222impl std::error::Error for ParamError {}
223
224pub const _DEFAULT_CUTOFFS_FOR_DOCS: (u32, u32) = (DEFAULT_CUTOFF_REFS, DEFAULT_CUTOFF_READS);
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229
230    #[test]
231    fn validates_current_cpp_defaults() {
232        let mut p = BuildParams::new(GraphInput::References, "out".to_string());
233        p.seqs.push("data/refs1.fa".to_string());
234        assert_eq!(p.cutoff(), 1);
235        assert!(p.validate().is_ok());
236    }
237
238    #[test]
239    fn rejects_even_k_and_zero_cutoff() {
240        let mut p = BuildParams::new(GraphInput::Reads, "out".to_string());
241        p.seqs.push("data/reads.fq".to_string());
242        p.k = 32;
243        assert_eq!(p.validate(), Err(ParamError::InvalidK(32)));
244        p.k = 31;
245        p.cutoff = Some(0);
246        assert_eq!(p.validate(), Err(ParamError::InvalidCutoff));
247    }
248
249    #[test]
250    fn explicit_memory_budget_bounds_phase_concurrency() {
251        let mut p = BuildParams::new(GraphInput::References, "out".to_string());
252        p.threads = 256;
253        p.max_memory_gb = Some(16);
254
255        assert_eq!(p.partition_workers(1_000), 256);
256        assert_eq!(p.local_workers(), 64);
257        assert_eq!(p.post_local_workers(), 128);
258        assert_eq!(p.partition_workers(17), 17);
259    }
260
261    #[test]
262    fn unrestricted_concurrency_follows_the_user_request() {
263        let mut p = BuildParams::new(GraphInput::References, "out".to_string());
264        p.threads = 96;
265
266        assert_eq!(p.partition_workers(1_000), 96);
267        assert_eq!(p.local_workers(), 96);
268        assert_eq!(p.post_local_workers(), 96);
269    }
270
271    #[test]
272    fn ample_memory_budget_preserves_non_power_of_two_request() {
273        let mut p = BuildParams::new(GraphInput::References, "out".to_string());
274        p.threads = 96;
275        p.max_memory_gb = Some(24);
276
277        assert_eq!(p.partition_workers(1_000), 96);
278        assert_eq!(p.local_workers(), 96);
279        assert_eq!(p.post_local_workers(), 96);
280    }
281
282    #[test]
283    fn binding_memory_budget_only_reduces_constrained_phases() {
284        let mut p = BuildParams::new(GraphInput::References, "out".to_string());
285        p.threads = 96;
286        p.max_memory_gb = Some(12);
287
288        assert_eq!(p.partition_workers(1_000), 96);
289        assert_eq!(p.local_workers(), 32);
290        assert_eq!(p.post_local_workers(), 96);
291    }
292
293    #[test]
294    fn rejects_zero_memory_budget() {
295        let mut p = BuildParams::new(GraphInput::References, "out".to_string());
296        p.seqs.push("data/refs1.fa".to_string());
297        p.max_memory_gb = Some(0);
298        assert_eq!(p.validate(), Err(ParamError::InvalidMemoryLimit));
299    }
300}