#![allow(clippy::module_name_repetitions)]
use super::num_cpus;
use std::time::Duration;
pub const DEFAULT_SHUFFLE_SEED: u64 = 0x517c_c1b7_2722_0a95;
pub const MAX_LOAD_FILE_SIZE: u64 = 256 * 1024 * 1024;
#[derive(Debug, Clone)]
pub struct PipelineConfig {
pub num_workers: usize,
pub prefetch_size: usize,
pub on_error: ErrorPolicy,
pub seed: Option<u64>,
pub epochs: usize,
pub channel_chunk_size: usize,
pub pending_sequence_limit: usize,
pub sequence_gap_timeout: Duration,
pub source_timeout: Option<Duration>,
pub drop_last: bool,
pub pin_threads: bool,
pub(crate) shard: Option<(usize, usize)>,
#[doc(hidden)]
pub(crate) test_start_sequence: u64,
}
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ErrorPolicy {
Skip,
Fail,
}
impl Default for PipelineConfig {
fn default() -> Self {
const FALLBACK_WORKERS: usize = 4;
let workers = num_cpus().unwrap_or_else(|error| {
tracing::warn!(
"tenshift: could not detect available parallelism ({error}); defaulting to {FALLBACK_WORKERS} workers. Override with PipelineConfig::workers()."
);
FALLBACK_WORKERS
});
Self {
num_workers: workers,
prefetch_size: workers.saturating_mul(2).max(8),
on_error: ErrorPolicy::Skip,
seed: None,
epochs: 1,
channel_chunk_size: 64,
pending_sequence_limit: 1000,
sequence_gap_timeout: Duration::from_secs(30),
source_timeout: None,
drop_last: false,
pin_threads: false,
shard: None,
test_start_sequence: 0,
}
}
}
impl PipelineConfig {
#[must_use]
pub fn prefetch_auto(mut self) -> Self {
self.prefetch_size = self.num_workers.saturating_mul(2).max(2);
self
}
}