dataloader-rs 0.2.0

use crate::sampler::Sampler;

/// Shards indices across `world_size` workers for distributed training.
pub struct DistributedSampler<S: Sampler> {
    inner: S,
    rank: usize,
    world_size: usize,
}

impl<S: Sampler> DistributedSampler<S> {
    /// Create a new `DistributedSampler`.
    ///
    /// # Panics
    ///
    /// Panics if `rank >= world_size` or `world_size == 0`.
    pub fn new(inner: S, rank: usize, world_size: usize) -> Self {
        assert!(world_size > 0, "world_size must be > 0");
        assert!(rank < world_size, "rank must be < world_size");
        Self {
            inner,
            rank,
            world_size,
        }
    }
}

impl<S: Sampler> Sampler for DistributedSampler<S> {
    fn indices(&mut self, dataset_len: usize) -> Vec<usize> {
        let all = self.inner.indices(dataset_len);
        // As in PyTorch: pad the epoch to a multiple of `world_size` by
        // wrapping around, so every rank gets as many indices, and give this
        // rank every `world_size`-th position from `rank`.
        (0..all.len().div_ceil(self.world_size))
            .map(|i| all[(self.rank + i * self.world_size) % all.len()])
            .collect()
    }

    fn len(&self, dataset_len: usize) -> usize {
        self.inner.len(dataset_len).div_ceil(self.world_size)
    }
}

#[cfg(test)]
mod tests {
    use crate::sampler::{DistributedSampler, RandomSampler, Sampler, SequentialSampler};

    #[test]
    fn distributed_partitions_have_expected_total_len() {
        let world_size = 3;
        let n = 10;
        let mut all_indices: Vec<usize> = Vec::new();
        for rank in 0..world_size {
            let mut ds = DistributedSampler::new(SequentialSampler, rank, world_size);
            all_indices.extend(ds.indices(n));
        }
        assert_eq!(all_indices.len(), 12);
    }

    #[test]
    fn distributed_equal_length_per_rank() {
        let world_size = 4;
        let n = 10;
        let lengths: Vec<usize> = (0..world_size)
            .map(|rank| {
                let mut ds = DistributedSampler::new(SequentialSampler, rank, world_size);
                ds.indices(n).len()
            })
            .collect();
        assert!(lengths.windows(2).all(|w| w[0] == w[1]));
        assert_eq!(lengths[0], 3);
    }

    #[test]
    fn distributed_len_matches_indices() {
        for (n, world_size) in [(10, 3), (10, 4), (8, 4), (1, 4), (0, 2)] {
            for rank in 0..world_size {
                let mut ds = DistributedSampler::new(SequentialSampler, rank, world_size);
                assert_eq!(ds.len(n), ds.indices(n).len(), "n={n} world={world_size}");
            }
        }
    }

    #[test]
    fn distributed_pads_by_wrapping_around() {
        let shard = |n, rank, world_size| {
            DistributedSampler::new(SequentialSampler, rank, world_size).indices(n)
        };
        // 10 indices padded to 12 with 0 and 1.
        assert_eq!(shard(10, 0, 4), vec![0, 4, 8]);
        assert_eq!(shard(10, 2, 4), vec![2, 6, 0]);
        assert_eq!(shard(10, 3, 4), vec![3, 7, 1]);
        // Padding longer than the epoch wraps more than once.
        for rank in 0..4 {
            assert_eq!(shard(1, rank, 4), vec![0]);
        }
        assert_eq!(shard(0, 1, 4), Vec::<usize>::new());
    }

    #[test]
    fn distributed_exact_divisor() {
        let world_size = 4;
        let n = 8;
        for rank in 0..world_size {
            let mut ds = DistributedSampler::new(SequentialSampler, rank, world_size);
            assert_eq!(ds.indices(n).len(), 2);
        }
    }

    #[test]
    #[should_panic]
    fn distributed_rank_out_of_bounds_panics() {
        DistributedSampler::new(SequentialSampler, 4, 4);
    }

    #[test]
    #[should_panic]
    fn distributed_zero_world_size_panics() {
        DistributedSampler::new(SequentialSampler, 0, 0);
    }

    #[test]
    fn distributed_wraps_random_sampler() {
        let world_size = 2;
        let n = 6;
        let mut ds = DistributedSampler::new(RandomSampler::new(99), 0, world_size);
        let indices = ds.indices(n);
        assert_eq!(indices.len(), n / world_size);
        assert!(indices.iter().all(|&i| i < n));
    }
}