Fix double sharding in ShuffleBuffer

monai/data/grid_dataset.py

@@ -218,6 +218,8 @@ class GridPatchDataset(IterableDataset):
 
     """
 
+    _shards_by_worker = True
+
     def __init__(
         self,
         data: Iterable | Sequence,
@@ -404,6 +406,8 @@ class PatchDataset(IterableDataset):
 
     """
 
+    _shards_by_worker = True
+
     def __init__(
         self, data: Sequence, patch_func: Callable, samples_per_image: int = 1, transform: Callable | None = None
     ) -> None:

monai/data/iterable_dataset.py

@@ -25,6 +25,23 @@
 pd, _ = optional_import("pandas")
 
 
+def _source_shards_by_worker(data: Iterable[Any]) -> bool:
+    """Return whether the source declares that its iterator partitions by worker.
+
+    Args:
+        data: iterable source to inspect through its type's method resolution order.
+
+    Returns:
+        ``True`` if the first class defining ``__iter__`` declares a truthy
+        ``_shards_by_worker`` value on itself, or ``False`` if the declaration
+        is missing or false.
+    """
+    for source_type in type(data).__mro__:
+        if "__iter__" in source_type.__dict__:
+            return bool(source_type.__dict__.get("_shards_by_worker", False))
+    return False
+
+
 class IterableDataset(_TorchIterableDataset):
     """
     A generic dataset for iterable data source and an optional callable data transform
@@ -40,6 +57,8 @@ class IterableDataset(_TorchIterableDataset):
 
     """
 
+    _shards_by_worker = True
+
     def __init__(self, data: Iterable[Any], transform: Callable | None = None) -> None:
         """
         Args:
@@ -75,6 +94,13 @@ class ShuffleBuffer(Randomizable, IterableDataset):
             every iter() call, refer to the PyTorch idea:
             https://github.com/pytorch/pytorch/blob/v1.10.0/torch/utils/data/distributed.py#L98.
         epochs: number of epochs to iterate over the dataset, default to 1, -1 means infinite epochs.
+        source_shards_by_worker: whether ``data`` already partitions its stream
+            using ``torch.utils.data.get_worker_info``. ``None`` automatically
+            recognizes built-in MONAI sources that declare worker partitioning.
+            A subclass that overrides iteration without declaring that capability
+            is treated as unsharded. ``True`` avoids a second worker partition
+            for any worker-aware source, and ``False`` preserves the outer
+            partition for unsharded iterable datasets.
 
     Note:
         Both ``monai.data.DataLoader`` and ``torch.utils.data.DataLoader`` do not seed this class (as a subclass of
@@ -97,11 +123,37 @@ def run():
 
     """
 
-    def __init__(self, data, transform=None, buffer_size: int = 512, seed: int = 0, epochs: int = 1) -> None:
+    _shards_by_worker = True
+
+    def __init__(
+        self,
+        data,
+        transform=None,
+        buffer_size: int = 512,
+        seed: int = 0,
+        epochs: int = 1,
+        source_shards_by_worker: bool | None = None,
+    ) -> None:
+        """Initialize the shuffle buffer.
+
+        Args:
+            data: input data source to load, shuffle, and optionally transform.
+            transform: a callable data transform applied to each yielded item.
+            buffer_size: maximum number of items stored before random popping.
+            seed: random seed used to initialize the worker random states.
+            epochs: number of source iterations, where ``-1`` means infinite.
+            source_shards_by_worker: whether ``data`` already partitions its
+                stream using ``torch.utils.data.get_worker_info``. ``None``
+                automatically recognizes built-in MONAI sources that declare
+                worker partitioning.
+        """
         super().__init__(data=data, transform=transform)
         self.size = buffer_size
         self.seed = seed
         self.epochs = epochs
+        self.source_shards_by_worker = (
+            _source_shards_by_worker(data) if source_shards_by_worker is None else source_shards_by_worker
+        )
         self._idx = 0
 
     def randomized_pop(self, buffer):
@@ -122,14 +174,24 @@ def generate_item(self):
             yield self.randomized_pop(buffer)
 
     def __iter__(self):
-        """
-        Randomly pop buffered items from `self.data`.
-        Multiple dataloader workers sharing this dataset will generate identical item sequences.
+        """Randomly pop buffered items from ``self.data``.
+
+        Yields:
+            Items from the shuffled source after applying the optional transform.
+
+        Raises:
+            RuntimeError: When the optional transform raises an exception.
         """
         self.seed += 1
         super().set_random_state(seed=self.seed)  # make all workers in sync
         for _ in range(self.epochs) if self.epochs >= 0 else iter(int, 1):
-            yield from IterableDataset(self.generate_item(), transform=self.transform)
+            if self.source_shards_by_worker:
+                for item in self.generate_item():
+                    if self.transform is not None:
+                        item = apply_transform(self.transform, item)
+                    yield item
+            else:
+                yield from IterableDataset(self.generate_item(), transform=self.transform)
 
     def randomize(self, size: int) -> None:
         self._idx = self.R.randint(size)
@@ -197,6 +259,8 @@ class CSVIterableDataset(IterableDataset):
 
     """
 
+    _shards_by_worker = True
+
     def __init__(
         self,
         src: str | Sequence[str] | Iterable | Sequence[Iterable],
@@ -278,4 +342,5 @@ def __iter__(self):
                 data=self._flattened(), transform=self.transform, buffer_size=self.buffer_size, seed=self.seed
             )
             yield from buffer
-        yield from IterableDataset(data=self._flattened(), transform=self.transform)
+        else:
+            yield from IterableDataset(data=self._flattened(), transform=self.transform)

tests/data/test_csv_iterable_dataset.py

@@ -24,8 +24,39 @@
 from tests.test_utils import skip_if_windows
 
 
+class _OneShotChunks:
+    """Yield one DataFrame chunk and reject a second iteration."""
+
+    def __init__(self):
+        self.iterations = 0
+
+    def __iter__(self):
+        """Yield the source's single DataFrame chunk.
+
+        Yields:
+            A DataFrame containing two test records.
+
+        Raises:
+            RuntimeError: When the source is iterated more than once.
+        """
+        self.iterations += 1
+        if self.iterations > 1:
+            raise RuntimeError("one-shot source reused")
+        yield pd.DataFrame({"subject_id": ["s0", "s1"], "label": [0, 1]})
+
+
 @skip_if_windows
 class TestCSVIterableDataset(unittest.TestCase):
+    def test_shuffle_consumes_one_shot_source_once(self):
+        source = _OneShotChunks()
+        dataset = CSVIterableDataset(src=source, chunksize=2, buffer_size=2, shuffle=True, seed=7)
+
+        items = list(dataset)
+
+        self.assertEqual(source.iterations, 1)
+        self.assertEqual(len(items), 2)
+        self.assertEqual({item["subject_id"] for item in items}, {"s0", "s1"})
+
     def test_values(self):
         with tempfile.TemporaryDirectory() as tempdir:
             test_data1 = [

tests/data/test_shuffle_buffer.py

@@ -13,13 +13,47 @@
 
 import sys
 import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
 
 import numpy as np
+from torch.utils.data import IterableDataset as TorchIterableDataset
 
-from monai.data import DataLoader, ShuffleBuffer
+from monai.data import DataLoader, GridPatchDataset, IterableDataset, PatchIter, ShuffleBuffer
+from monai.data import iterable_dataset as iterable_dataset_module
 from monai.utils import convert_data_type
 
 
+class _UnshardedMonaiIterable(IterableDataset):
+    """MONAI iterable subclass that intentionally does not partition itself."""
+
+    def __iter__(self):
+        """Yield every item from the unsharded source.
+
+        Yields:
+            Items from ``self.data``.
+        """
+        yield from self.data
+
+
+class _WorkerShardedTorchIterable(TorchIterableDataset):
+    """PyTorch iterable source that partitions itself across workers."""
+
+    def __init__(self, size):
+        self.size = size
+
+    def __iter__(self):
+        """Yield the integer indices assigned to the current worker.
+
+        Yields:
+            Integer indices from the worker's partition of ``range(self.size)``.
+        """
+        worker_info = iterable_dataset_module.get_worker_info()
+        num_workers = worker_info.num_workers if worker_info is not None else 1
+        worker_id = worker_info.id if worker_info is not None else 0
+        yield from range(worker_id, self.size, num_workers)
+
+
 class TestShuffleBuffer(unittest.TestCase):
     def test_shape(self):
         buffer = ShuffleBuffer([1, 2, 3, 4], seed=0)
@@ -37,6 +71,72 @@ def test_shape(self):
             np.testing.assert_allclose(output, [[2, 3], [1, 4]], err_msg=f"seed {buffer.seed}")
             np.testing.assert_allclose(output2, [[1, 4], [2, 3]], err_msg=f"seed {buffer.seed}")
 
+    def test_monai_iterable_source_is_detected_as_worker_sharded(self):
+        """Verify MONAI iterable sources avoid a second worker partition by default."""
+        outputs = []
+        for worker_id in range(2):
+            source = IterableDataset(range(40))
+            buffer = ShuffleBuffer(source, buffer_size=8, seed=7)
+            worker_info = SimpleNamespace(num_workers=2, id=worker_id)
+            with patch("monai.data.iterable_dataset.get_worker_info", return_value=worker_info):
+                outputs.extend(buffer)
+
+        self.assertEqual(len(outputs), 40)
+        self.assertEqual(set(outputs), set(range(40)))
+
+    def test_worker_sharded_source_is_not_sharded_twice(self):
+        """Verify an explicitly worker-sharded source is not repartitioned."""
+        sources = [IterableDataset(range(40)), _WorkerShardedTorchIterable(40)]
+        for source in sources:
+            outputs = []
+            for worker_id in range(2):
+                buffer = ShuffleBuffer(
+                    source, transform=lambda item: item + 40, buffer_size=8, seed=7, source_shards_by_worker=True
+                )
+                worker_info = SimpleNamespace(num_workers=2, id=worker_id)
+                with patch("monai.data.iterable_dataset.get_worker_info", return_value=worker_info):
+                    outputs.extend(buffer)
+
+            self.assertEqual(len(outputs), 40)
+            self.assertEqual(set(outputs), set(range(40, 80)))
+
+    def test_grid_patch_dataset_is_not_sharded_twice(self):
+        """Verify every grid patch is yielded exactly once across workers."""
+        images = [np.arange(16).reshape(1, 4, 4), np.arange(16, 32).reshape(1, 4, 4)]
+        expected_patches = [
+            (0, 1, 4, 5),
+            (2, 3, 6, 7),
+            (8, 9, 12, 13),
+            (10, 11, 14, 15),
+            (16, 17, 20, 21),
+            (18, 19, 22, 23),
+            (24, 25, 28, 29),
+            (26, 27, 30, 31),
+        ]
+        outputs = []
+        for worker_id in range(2):
+            source = GridPatchDataset(data=images, patch_iter=PatchIter(patch_size=(2, 2)), with_coordinates=False)
+            buffer = ShuffleBuffer(source, buffer_size=3, seed=7)
+            worker_info = SimpleNamespace(num_workers=2, id=worker_id)
+            with patch("monai.data.iterable_dataset.get_worker_info", return_value=worker_info):
+                outputs.extend(tuple(item.ravel().tolist()) for item in buffer)
+
+        self.assertCountEqual(outputs, expected_patches)
+
+    def test_unsharded_monai_subclass_keeps_outer_worker_partition(self):
+        """Verify default and explicit unsharded modes preserve outer partitioning."""
+        for source_shards_by_worker in (None, False):
+            outputs = []
+            for worker_id in range(2):
+                source = _UnshardedMonaiIterable(range(40))
+                buffer = ShuffleBuffer(source, buffer_size=8, seed=7, source_shards_by_worker=source_shards_by_worker)
+                worker_info = SimpleNamespace(num_workers=2, id=worker_id)
+                with patch("monai.data.iterable_dataset.get_worker_info", return_value=worker_info):
+                    outputs.extend(buffer)
+
+            self.assertEqual(len(outputs), 40)
+            self.assertEqual(set(outputs), set(range(40)))
+
     def test_epochs(self):
         buffer = ShuffleBuffer([1, 2, 3, 4], seed=0, epochs=2)
         output = [convert_data_type(x, np.ndarray)[0] for x in DataLoader(dataset=buffer, batch_size=2)]