dataloader_rs/loader/source/
sync.rs

1use std::sync::Arc;
2
3use super::{LoaderSource, sealed::Sealed, settings::EpochSettings, threads::pool_threads};
4use crate::{collator::Collator, dataset::Dataset, loader::sync::SyncEpoch, transform::Transform};
5
6/// A [`Dataset`], as a [`DataLoader`](crate::DataLoader) built by
7/// [`DataLoader::builder`](crate::DataLoader::builder) holds it.
8pub struct SyncSource<D>(Arc<D>);
9
10impl<D> SyncSource<D> {
11    pub(in crate::loader) fn new(dataset: D) -> Self {
12        Self(Arc::new(dataset))
13    }
14}
15
16impl<D: Dataset> LoaderSource for SyncSource<D> {
17    type Dataset = D;
18
19    fn dataset(&self) -> &D {
20        &self.0
21    }
22}
23
24impl<D: Dataset> Sealed for SyncSource<D> {
25    type Item = D::Item;
26    type Epoch<T, C, O>
27        = SyncEpoch<D, T, C>
28    where
29        T: Transform<D::Item, Output = O>,
30        C: Collator<O>;
31
32    fn len(&self) -> usize {
33        self.0.len()
34    }
35
36    /// One pool for batch-level and item-level work: `num_workers` batches, each with
37    /// `item_concurrency` items in flight; none when both are at their defaults.
38    fn pool_threads(
39        &self,
40        num_workers: usize,
41        item_concurrency: usize,
42        _transform_is_identity: bool,
43    ) -> Option<usize> {
44        (num_workers > 0 || item_concurrency > 1)
45            .then(|| pool_threads(num_workers, item_concurrency))
46    }
47
48    fn epoch<T, C, O>(
49        &self,
50        transform: &Arc<T>,
51        collator: &Arc<C>,
52        pool: Option<&Arc<rayon::ThreadPool>>,
53        settings: EpochSettings,
54        chunks: Vec<Vec<usize>>,
55    ) -> SyncEpoch<D, T, C>
56    where
57        T: Transform<D::Item, Output = O>,
58        C: Collator<O>,
59    {
60        SyncEpoch::start(&self.0, transform, collator, pool, settings, chunks)
61    }
62}
dataloader_rs/loader/source/sync.rs

dataloader_rs/loader/source/
sync.rs
