dataloader_rs/loader/
builder.rs

1use std::sync::Arc;
2
3#[cfg(feature = "torch-rs")]
4use crate::collator::TorchPinnedCollator;
5#[cfg(feature = "async")]
6use crate::loader::source::AsyncSource;
7use crate::{
8    collator::{Collator, VecCollator},
9    loader::{
10        dataloader::DataLoader,
11        pool::{LoaderPool, ThreadMain},
12        source::LoaderSource,
13    },
14    sampler::{BatchSampler, Sampler, SequentialSampler},
15    transform::{Identity, Transform},
16};
17#[cfg(feature = "async")]
18use tokio::runtime::Handle;
19
20/// Builder for [`DataLoader`].
21pub struct DataLoaderBuilder<Src, S, C, T = Identity> {
22    source: Src,
23    sampler: S,
24    transform: T,
25    collator: C,
26    settings: Settings,
27}
28
29/// The settings that do not depend on the builder's type parameters, carried
30/// over as they are when [`sampler`](DataLoaderBuilder::sampler),
31/// [`transform`](DataLoaderBuilder::transform),
32/// [`collator`](DataLoaderBuilder::collator) or `pin_memory` change a type.
33struct Settings {
34    batch_size: usize,
35    drop_last: bool,
36    prefetch_depth: usize,
37    /// Batches processed in parallel.
38    num_workers: usize,
39    /// Items of one batch fetched in parallel.
40    item_concurrency: usize,
41    /// Main function of the pool's threads; `None` runs rayon's loop as is.
42    thread_main: Option<ThreadMain>,
43}
44
45impl<Src: LoaderSource> DataLoaderBuilder<Src, SequentialSampler, VecCollator> {
46    pub(super) fn new(source: Src) -> Self {
47        DataLoaderBuilder {
48            source,
49            sampler: SequentialSampler,
50            transform: Identity,
51            collator: VecCollator,
52            settings: Settings {
53                batch_size: 1,
54                drop_last: false,
55                prefetch_depth: 2,
56                num_workers: 0,
57                item_concurrency: 1,
58                thread_main: None,
59            },
60        }
61    }
62}
63
64impl<Src, S, C, T> DataLoaderBuilder<Src, S, C, T> {
65    /// Number of items per batch. Default: `1`.
66    pub fn batch_size(mut self, n: usize) -> Self {
67        assert!(n > 0, "batch_size must be > 0");
68        self.settings.batch_size = n;
69        self
70    }
71
72    /// Number of batches workers may run ahead of the consumer. At most
73    /// `num_workers + prefetch_depth` batches are in flight or buffered.
74    /// Default: `2`.
75    pub fn prefetch_depth(mut self, n: usize) -> Self {
76        assert!(n > 0, "prefetch_depth must be > 0");
77        self.settings.prefetch_depth = n;
78        self
79    }
80
81    /// Whether to drop a non-full final batch. Default: `false`.
82    pub fn drop_last(mut self, v: bool) -> Self {
83        self.settings.drop_last = v;
84        self
85    }
86
87    /// Number of batches processed in parallel — PyTorch's `num_workers`.
88    ///
89    /// Batches run on the loader's thread pool, built once and reused by every
90    /// epoch. `0` (default) processes batches directly in `Iterator::next`.
91    pub fn num_workers(mut self, n: usize) -> Self {
92        self.settings.num_workers = n;
93        self
94    }
95
96    /// Number of items of one batch fetched in parallel. Default: `1`.
97    ///
98    /// With `1`, a batch is fetched with one
99    /// [`Dataset::get_batch`](crate::Dataset::get_batch) call. With `n > 1`,
100    /// each batch is fetched by `n` concurrent lanes, each calling
101    /// [`Dataset::get`](crate::Dataset::get) for the batch's next item, so at
102    /// most `n` items of a batch are in flight. Useful for I/O-bound datasets (network, object
103    /// storage) and for few, large CPU-bound batches.
104    ///
105    /// The loader's pool has `max(num_workers, 1) × item_concurrency`
106    /// threads; sizing it for the machine is the caller's choice.
107    pub fn item_concurrency(mut self, n: usize) -> Self {
108        assert!(n > 0, "item_concurrency must be > 0");
109        self.settings.item_concurrency = n;
110        self
111    }
112
113    /// Former name of [`item_concurrency`](Self::item_concurrency).
114    #[deprecated(since = "0.3.0", note = "renamed to `item_concurrency`")]
115    pub fn items_in_parallel(self, n: usize) -> Self {
116        self.item_concurrency(n)
117    }
118
119    /// Run each pool thread through `main`, which must call
120    /// `ThreadBuilder::run` (the Python binding keeps the thread attached to
121    /// the interpreter for its whole life).
122    #[cfg(feature = "python")]
123    pub(crate) fn thread_main(mut self, main: ThreadMain) -> Self {
124        self.settings.thread_main = Some(main);
125        self
126    }
127
128    /// Pin Tensor batches in host memory before they are sent to the consumer.
129    ///
130    /// Available only for `tch::Tensor` batches under `torch-rs` feature.
131    #[cfg(feature = "torch-rs")]
132    pub fn pin_memory(self, enabled: bool) -> DataLoaderBuilder<Src, S, TorchPinnedCollator<C>, T>
133    where
134        Src: LoaderSource,
135        T: Transform<Src::Item>,
136        C: Collator<T::Output, Batch = tch::Tensor>,
137    {
138        DataLoaderBuilder {
139            source: self.source,
140            sampler: self.sampler,
141            transform: self.transform,
142            collator: TorchPinnedCollator::new(self.collator, enabled),
143            settings: self.settings,
144        }
145    }
146
147    /// Replace the sampler and preserve all other settings.
148    pub fn sampler<S2: Sampler>(self, sampler: S2) -> DataLoaderBuilder<Src, S2, C, T> {
149        DataLoaderBuilder {
150            source: self.source,
151            sampler,
152            transform: self.transform,
153            collator: self.collator,
154            settings: self.settings,
155        }
156    }
157
158    /// Apply `transform` to every item, with its index, before collation; it
159    /// replaces any previous one (chain them with
160    /// [`Transform::then`]). Default: [`Identity`], at no cost.
161    ///
162    /// It runs where the item is loaded: right after `get` / `get_batch` on the
163    /// worker for a [`Dataset`](crate::Dataset) (in each item lane when
164    /// [`item_concurrency`](Self::item_concurrency) is above 1, so items are
165    /// transformed in parallel), or on the loader's pool once the item is
166    /// fetched for an `AsyncDataset`, so the async runtime never does CPU work.
167    /// The collator receives the transformed items. An error fails the item's
168    /// batch; a panic is re-raised by the iterator.
169    ///
170    /// ```rust
171    /// use dataloader_rs::{DataLoader, Dataset, error::Result};
172    ///
173    /// struct Numbers;
174    ///
175    /// impl Dataset for Numbers {
176    ///     type Item = u64;
177    ///     fn get(&self, index: usize) -> Result<u64> { Ok(index as u64) }
178    ///     fn len(&self) -> usize { 4 }
179    /// }
180    ///
181    /// let mut loader = DataLoader::builder(Numbers)
182    ///     .batch_size(2)
183    ///     .transform(|index: usize, x: u64| -> Result<String> { Ok(format!("{index}:{}", x * 10)) })
184    ///     .build();
185    /// let batches: Vec<Vec<String>> = loader.iter().map(|b| b.unwrap()).collect();
186    /// assert_eq!(batches, [["0:0", "1:10"], ["2:20", "3:30"]]);
187    /// ```
188    pub fn transform<T2>(self, transform: T2) -> DataLoaderBuilder<Src, S, C, T2>
189    where
190        Src: LoaderSource,
191        T2: Transform<Src::Item>,
192    {
193        DataLoaderBuilder {
194            source: self.source,
195            sampler: self.sampler,
196            transform,
197            collator: self.collator,
198            settings: self.settings,
199        }
200    }
201
202    /// Replace the collator and preserve all other settings.
203    pub fn collator<C2>(self, collator: C2) -> DataLoaderBuilder<Src, S, C2, T> {
204        DataLoaderBuilder {
205            source: self.source,
206            sampler: self.sampler,
207            transform: self.transform,
208            collator,
209            settings: self.settings,
210        }
211    }
212
213    /// Finalize configuration and construct a [`DataLoader`].
214    pub fn build(self) -> DataLoader<Src, S, C, T>
215    where
216        Src: LoaderSource,
217        S: Sampler,
218        T: Transform<Src::Item>,
219        C: Collator<T::Output>,
220    {
221        let Settings {
222            batch_size,
223            drop_last,
224            prefetch_depth,
225            num_workers,
226            item_concurrency,
227            thread_main,
228        } = self.settings;
229        self.source.start();
230        let pool = self
231            .source
232            .pool_threads(num_workers, item_concurrency, self.transform.is_identity())
233            .map(|threads| LoaderPool::new(threads, thread_main));
234
235        DataLoader {
236            source: self.source,
237            batch_sampler: BatchSampler::new(self.sampler, batch_size, drop_last),
238            transform: Arc::new(self.transform),
239            collator: Arc::new(self.collator),
240            prefetch_depth,
241            num_workers,
242            item_concurrency,
243            pool,
244        }
245    }
246}
247
248#[cfg(feature = "async")]
249impl<A, S, C, T> DataLoaderBuilder<AsyncSource<A>, S, C, T> {
250    /// Fetch items on the caller's tokio runtime instead of the loader's own. It must be a
251    /// multi-threaded runtime: the loader's iterator blocks while it waits for a batch.
252    ///
253    /// # Panics
254    /// If [`runtime_threads`](Self::runtime_threads) is set too.
255    pub fn runtime(mut self, handle: Handle) -> Self {
256        self.source.runtime.set_handle(handle);
257        self
258    }
259
260    /// Worker threads of the loader's own tokio runtime, built by [`build`](Self::build).
261    /// Default: the CPUs available to the process
262    /// ([`available_parallelism`](std::thread::available_parallelism): affinity and cgroup
263    /// quota included; 1 if unknown), at most [`MAX_RUNTIME_THREADS`](crate::MAX_RUNTIME_THREADS).
264    ///
265    /// These threads receive responses, run TLS and copy bodies; the number of requests in
266    /// flight is `max(num_workers, 1) × item_concurrency`, and transforms and collation run on
267    /// rayon. One thread can drive many requests in flight.
268    ///
269    /// # Panics
270    /// If `n` is 0 or above [`MAX_RUNTIME_THREADS`](crate::MAX_RUNTIME_THREADS), or if [`runtime`](Self::runtime) is set
271    /// too.
272    pub fn runtime_threads(mut self, n: usize) -> Self {
273        self.source.runtime.set_threads(n);
274        self
275    }
276}
277
278#[cfg(test)]
279mod tests {
280    use crate::{
281        collator::Collator, dataset::Dataset, error::Result, loader::DataLoader, sampler::Sampler,
282    };
283
284    struct TinyDs(usize);
285
286    impl Dataset for TinyDs {
287        type Item = usize;
288
289        fn get(&self, index: usize) -> Result<Self::Item> {
290            Ok(index)
291        }
292
293        fn len(&self) -> usize {
294            self.0
295        }
296    }
297
298    struct ReverseSampler;
299
300    impl Sampler for ReverseSampler {
301        fn indices(&mut self, dataset_len: usize) -> Vec<usize> {
302            (0..dataset_len).rev().collect()
303        }
304    }
305
306    struct SumCollator;
307
308    impl Collator<usize> for SumCollator {
309        type Batch = usize;
310
311        fn collate(&self, items: Vec<usize>) -> Result<Self::Batch> {
312            Ok(items.into_iter().sum())
313        }
314    }
315
316    #[test]
317    fn builder_defaults_create_single_item_batches() {
318        let mut loader = DataLoader::builder(TinyDs(4)).build();
319        let batches: Vec<Vec<usize>> = loader.iter().map(|b| b.unwrap()).collect();
320        assert_eq!(batches, vec![vec![0], vec![1], vec![2], vec![3]]);
321    }
322
323    #[test]
324    fn builder_replaces_sampler_and_collator() {
325        let mut loader = DataLoader::builder(TinyDs(4))
326            .batch_size(2)
327            .sampler(ReverseSampler)
328            .collator(SumCollator)
329            .build();
330
331        let sums: Vec<usize> = loader.iter().map(|b| b.unwrap()).collect();
332        assert_eq!(sums, vec![5, 1]); // [3,2] and [1,0]
333    }
334
335    #[test]
336    #[should_panic]
337    fn builder_rejects_zero_batch_size() {
338        let _ = DataLoader::builder(TinyDs(4)).batch_size(0);
339    }
340
341    #[test]
342    #[should_panic]
343    fn builder_rejects_zero_prefetch_depth() {
344        let _ = DataLoader::builder(TinyDs(4)).prefetch_depth(0);
345    }
346}
dataloader_rs/loader/builder.rs

dataloader_rs/loader/
builder.rs
