1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
//! Builder logic for assembling a tenshift pipeline.
//!
//! This module is the front door to the pipeline architecture: it captures the
//! source, stateless worker stages, collector stages, and runtime configuration
//! before execution is handed off to the threaded executor.
#![allow(clippy::module_name_repetitions)]
use super::{CollateMode, ErrorPolicy, PipelineConfig, PipelineIterator, Stage};
use crate::error::Result;
use crate::sample::Sample;
use crate::source::Source;
use crate::transform::{FilterTransform, FlatMapTransform, MapTransform};
use std::sync::Arc;
use std::time::Duration;
/// A composable data loading pipeline.
///
/// Build it with [`Pipeline::from_source`], chain transforms, then iterate.
pub struct Pipeline {
pub(crate) source: Box<dyn Source>,
pub(crate) stages: Vec<Stage>,
pub(crate) config: PipelineConfig,
pub(crate) collate_mode: CollateMode,
}
impl Pipeline {
/// Create a pipeline from a data source.
pub fn from_source(source: impl Source + 'static) -> Self {
Self {
source: Box::new(source),
stages: Vec::new(),
config: PipelineConfig::default(),
collate_mode: CollateMode::Disabled,
}
}
/// Create a pipeline from an async data source.
///
/// This wraps the async source in a dedicated Tokio single-threaded runtime
/// that drives the async iterator natively without blocking the main event loops.
pub fn from_async_source(source: impl crate::source::AsyncSource + 'static) -> Self {
Self::from_source(crate::source::AsyncToSyncAdapter::new(source))
}
/// Set the number of parallel worker threads (default: number of CPUs, max 8).
pub fn workers(mut self, n: usize) -> Self {
self.config.num_workers = n.max(1);
self
}
/// Set the prefetch buffer size (number of items held ready).
pub fn prefetch(mut self, n: usize) -> Self {
self.config.prefetch_size = n.max(1);
self
}
/// Set the error handling policy.
pub fn on_error(mut self, policy: ErrorPolicy) -> Self {
self.config.on_error = policy;
self
}
/// Automatically set prefetch size to workers * 2.
///
/// Callers can use `.prefetch_auto()` instead of `.prefetch(2)` for CPU-bound pipelines.
pub fn prefetch_auto(mut self) -> Self {
self.config.prefetch_size = self.config.num_workers.saturating_mul(2).max(2);
self
}
/// Set prefetch count based on a memory budget and estimated sample size.
///
/// This avoids OOM when samples are large. For 100 MB samples with a
/// 256 MB budget: `prefetch_bytes(256 * 1024 * 1024, 100 * 1024 * 1024)`
/// yields `prefetch_size = 2`.
///
/// Both arguments are clamped to produce at least 1 prefetch slot.
pub fn prefetch_bytes(mut self, budget_bytes: usize, estimated_sample_bytes: usize) -> Self {
const MAX_PREFETCH_SIZE: usize = 10_000;
let sample_size = estimated_sample_bytes.max(1);
self.config.prefetch_size = (budget_bytes / sample_size).clamp(1, MAX_PREFETCH_SIZE);
self
}
/// Set a seed for deterministic shuffling.
pub fn seed(mut self, seed: u64) -> Self {
self.config.seed = Some(seed);
self
}
/// Set the number of epochs (passes over the data). 0 = infinite.
pub fn epochs(mut self, n: usize) -> Self {
self.config.epochs = n;
self
}
/// Set the internal channel chunk size (default: 16).
///
/// Higher values reduce synchronization overhead but use more memory.
pub fn chunk_size(mut self, n: usize) -> Self {
self.config.channel_chunk_size = n.max(1);
self
}
/// Set the maximum number of out-of-order processed chunks buffered in the collector.
pub fn pending_sequence_limit(mut self, n: usize) -> Self {
self.config.pending_sequence_limit = n.max(1);
self
}
/// Set how long the collector waits for a missing processed sequence before skipping it.
pub fn sequence_gap_timeout(mut self, timeout: Duration) -> Self {
self.config.sequence_gap_timeout = timeout;
self
}
/// Set how long the source may remain silent before the pipeline shuts down.
pub fn source_timeout(mut self, timeout: Duration) -> Self {
self.config.source_timeout = Some(timeout);
self
}
/// Control whether the final incomplete batch is emitted or discarded.
pub fn drop_last(mut self, enabled: bool) -> Self {
self.config.drop_last = enabled;
self
}
/// Pin worker threads to specific physical CPU cores to eliminate OS scheduler migration.
pub fn pin_threads(mut self, enabled: bool) -> Self {
self.config.pin_threads = enabled;
self
}
/// Shard the source across distributed ranks.
pub fn shard(mut self, rank: usize, world_size: usize) -> Self {
self.config.shard = Some((rank, world_size));
self
}
/// Add a map transform - apply a function to each sample.
pub fn map<F>(mut self, f: F) -> Self
where
F: Fn(Sample) -> Result<Sample> + Send + Sync + 'static,
{
self.stages
.push(Stage::Stateless(Box::new(MapTransform::new(f))));
self
}
/// Add a filter transform - keep only samples that match.
pub fn filter<F>(mut self, f: F) -> Self
where
F: Fn(&Sample) -> bool + Send + Sync + 'static,
{
self.stages
.push(Stage::Stateless(Box::new(FilterTransform::new(f))));
self
}
/// Add a `flat_map` transform - map one sample to zero or more output samples.
///
/// This is crucial for expanding single inputs into multiple training examples,
/// such as tokenizing long documents into sliding context windows for LLMs.
pub fn flat_map<F>(mut self, f: F) -> Self
where
F: Fn(Sample) -> Result<Vec<Sample>> + Send + Sync + 'static,
{
self.stages
.push(Stage::Stateless(Box::new(FlatMapTransform::new(f))));
self
}
/// Add a shuffle stage with the given buffer size.
pub fn shuffle(mut self, buffer_size: usize) -> Self {
self.stages.push(Stage::Shuffle(buffer_size));
self
}
/// Add a batching stage and enable the default collate function.
///
/// The default collate stacks matching tensor fields into a leading batch
/// dimension and emits one collated [`Sample`] per output batch.
pub fn batch(mut self, batch_size: usize) -> Self {
self.stages.push(Stage::Batch(batch_size));
if matches!(self.collate_mode, CollateMode::Disabled) {
self.collate_mode = CollateMode::Default;
}
self
}
/// Add a batching stage that preserves the raw `Vec<Sample>` output.
///
/// Use this when migrating code which expects batches to remain as
/// uncollated sample vectors.
pub fn batch_raw(mut self, batch_size: usize) -> Self {
self.stages.push(Stage::Batch(batch_size));
self
}
/// Run a custom collate function in the collector thread after batching.
///
/// ```rust
/// use tenshift_core::sample::Sample;
/// use tenshift_core::Pipeline;
/// use tenshift_core::sources::MemorySource;
///
/// let pipeline = Pipeline::from_source(MemorySource::new("demo", Vec::<Sample>::new()))
/// .batch_raw(2)
/// .collate_fn(|batch| Ok(match batch.into_iter().next() { Some(v) => v, None => Default::default() }));
///
/// let _ = pipeline;
/// ```
pub fn collate_fn<F>(mut self, f: F) -> Self
where
F: Fn(Vec<Sample>) -> Result<Sample> + Send + Sync + 'static,
{
self.collate_mode = CollateMode::Custom(Arc::new(f));
self
}
/// Enable the built-in collate function explicitly.
pub fn default_collate(mut self) -> Self {
self.collate_mode = CollateMode::Default;
self
}
/// Expose the resolved pipeline configuration (for tests and introspection).
#[must_use]
pub fn config(&self) -> &PipelineConfig {
&self.config
}
/// Expose the active collate mode (for tests and introspection).
#[must_use]
pub fn collate_mode(&self) -> CollateMode {
self.collate_mode.clone()
}
/// Start the pipeline and return an iterator.
///
/// # Errors
///
/// Returns an error when worker threads cannot be spawned or when the
/// pipeline configuration is internally inconsistent, including invalid
/// batching and collector stage ordering.
pub fn start(self) -> Result<PipelineIterator> {
crate::pipeline::executor::start_pipeline(self)
}
}