tenshift_core/pipeline/config.rs
1//! Runtime configuration for tenshift pipeline execution.
2//!
3//! These settings control how the source thread, worker pool, and collector
4//! cooperate. The module sits at the boundary between the public builder API
5//! and the executor's concrete threading behavior.
6//!
7//! # Auto-Scaling Heuristics
8//!
9//! The default configuration automatically adapts to the execution environment:
10//!
11//! ## Worker Count
12//! **Default:** `num_cpus().min(8)`
13//!
14//! Worker threads execute stateless transforms in parallel. The default caps
15//! at 8 because:
16//! - Most ML preprocessing is memory-bandwidth-bound, not CPU-bound
17//! - Beyond 8 workers, contention on memory channels often reduces throughput
18//! - Users with truly CPU-bound transforms can override with `.workers(n)`
19//!
20//! ## Prefetch Size
21//! **Default:** `num_workers * 2`, minimum 8
22//!
23//! The prefetch buffer sits between workers and collector. Its size balances:
24//! - **Worker starvation prevention:** Workers produce chunks faster than the
25//! collector consumes them during bursts. Without sufficient prefetch,
26//! workers block on full channels and stall.
27//! - **Memory bounds:** Each prefetch slot holds one full chunk. A fixed large
28//! buffer would OOM on high-core-count systems.
29//!
30//! The `* 2` multiplier provides one full chunk of headroom per worker,
31//! ensuring continuous work even with slight consumer jitter.
32//!
33//! ## Channel Chunk Size
34//! **Default:** 64 samples
35//!
36//! Samples flow through internal channels in chunks of 64. This amortizes:
37//! - **Channel synchronization overhead:** Sending 64 samples costs ~same as 1
38//! - **Cache locality:** Sequential processing within chunks improves hit rates
39//! - **Backpressure timing:** 64-sample granularity keeps OOM prevention responsive
40
41#![allow(clippy::module_name_repetitions)]
42
43use super::num_cpus;
44use std::time::Duration;
45
46/// The default random seed used for deterministic operations (e.g. shuffling).
47pub const DEFAULT_SHUFFLE_SEED: u64 = 0x517c_c1b7_2722_0a95;
48
49/// Maximum file size to load (256MB).
50pub const MAX_LOAD_FILE_SIZE: u64 = 256 * 1024 * 1024;
51
52/// Configuration for the pipeline.
53///
54/// Created via [`PipelineConfig::default()`](Default) or modified through
55/// the [`Pipeline`](crate::Pipeline) builder methods.
56#[derive(Debug, Clone)]
57pub struct PipelineConfig {
58 /// Number of worker threads for parallel data loading and transforms.
59 pub num_workers: usize,
60 /// Size of the prefetch buffer (number of items held in memory).
61 ///
62 /// Default: `num_workers * 2` with a minimum of 8. See module-level docs
63 /// for the rationale behind this heuristic.
64 pub prefetch_size: usize,
65 /// What to do when a source item or transform fails.
66 pub on_error: ErrorPolicy,
67 /// Optional seed for deterministic shuffling.
68 pub seed: Option<u64>,
69 /// Number of epochs (passes over the data). 0 = infinite.
70 pub epochs: usize,
71 /// Number of samples to chunk together in internal channels.
72 ///
73 /// Default: 64. Larger chunks amortize synchronization overhead but use
74 /// more memory. See module-level docs for the rationale.
75 pub channel_chunk_size: usize,
76 /// Maximum number of out-of-order processed chunks buffered while waiting for a missing sequence.
77 pub pending_sequence_limit: usize,
78 /// Maximum time to wait for a missing processed sequence before skipping it.
79 pub sequence_gap_timeout: Duration,
80 /// Maximum time to wait for the source to produce a new item before shutting down.
81 pub source_timeout: Option<Duration>,
82 /// Whether to drop the final incomplete batch instead of flushing it.
83 pub drop_last: bool,
84 /// Pin worker threads to specific physical CPU cores to eliminate OS scheduler migration.
85 pub pin_threads: bool,
86 pub(crate) shard: Option<(usize, usize)>,
87 #[doc(hidden)]
88 pub(crate) test_start_sequence: u64,
89}
90
91/// What to do when a data loading error occurs.
92#[non_exhaustive]
93#[derive(Debug, Clone, Copy, PartialEq, Eq)]
94pub enum ErrorPolicy {
95 /// Skip the bad item, log a warning, continue.
96 Skip,
97 /// Stop the entire pipeline on the first error.
98 Fail,
99}
100
101impl Default for PipelineConfig {
102 fn default() -> Self {
103 const FALLBACK_WORKERS: usize = 4;
104 let workers = num_cpus().unwrap_or_else(|error| {
105 tracing::warn!(
106 "tenshift: could not detect available parallelism ({error}); defaulting to {FALLBACK_WORKERS} workers. Override with PipelineConfig::workers()."
107 );
108 FALLBACK_WORKERS
109 });
110 Self {
111 num_workers: workers,
112 // Scale prefetch with worker count to keep all workers fed.
113 // Old default of 8 caused stalls on 64+ core machines.
114 // See module-level docs for the full rationale.
115 prefetch_size: workers.saturating_mul(2).max(8),
116 on_error: ErrorPolicy::Skip,
117 seed: None,
118 epochs: 1,
119 // Larger chunks amortize channel synchronization overhead.
120 // 64 is the sweet spot for ML workloads (see module docs).
121 channel_chunk_size: 64,
122 pending_sequence_limit: 1000,
123 sequence_gap_timeout: Duration::from_secs(30),
124 source_timeout: None,
125 drop_last: false,
126 pin_threads: false,
127 shard: None,
128 test_start_sequence: 0,
129 }
130 }
131}
132
133impl PipelineConfig {
134 /// Automatically set prefetch size to workers * 2.
135 ///
136 /// Callers can use `.prefetch_auto()` instead of `.prefetch(2)` for CPU-bound pipelines.
137 ///
138 /// # Rationale
139 ///
140 /// This recalculates the prefetch size based on the current worker count,
141 /// ensuring the buffer scales appropriately when workers are customized.
142 #[must_use]
143 pub fn prefetch_auto(mut self) -> Self {
144 self.prefetch_size = self.num_workers.saturating_mul(2).max(2);
145 self
146 }
147}