fix(data): detect unsharded iterable overrides

Author: yinkev

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,14 @@
 pd, _ = optional_import("pandas")
 
 
+def _source_shards_by_worker(data: Iterable[Any]) -> bool:
+    """Return whether the source declares that its iterator partitions by worker."""
+    source_type = type(data)
+    if "__iter__" in source_type.__dict__:
+        return bool(source_type.__dict__.get("_shards_by_worker", False))
+    return bool(getattr(source_type, "_shards_by_worker", False))
+
+
 class IterableDataset(_TorchIterableDataset):
     """
     A generic dataset for iterable data source and an optional callable data transform
@@ -40,6 +48,8 @@ class IterableDataset(_TorchIterableDataset):
 
     """
 
+    _shards_by_worker = True
+
     def __init__(self, data: Iterable[Any], transform: Callable | None = None) -> None:
         """
         Args:
@@ -77,9 +87,11 @@ class ShuffleBuffer(Randomizable, IterableDataset):
         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 MONAI ``IterableDataset`` sources, ``True`` avoids a
-            second worker partition for any worker-aware source, and ``False``
-            preserves the outer partition for unsharded iterable datasets.
+            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
@@ -102,6 +114,8 @@ def run():
 
     """
 
+    _shards_by_worker = True
+
     def __init__(
         self,
         data,
@@ -121,14 +135,15 @@ def __init__(
             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 MONAI ``IterableDataset`` sources.
+                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 = (
-            isinstance(data, IterableDataset) if source_shards_by_worker is None else source_shards_by_worker
+            _source_shards_by_worker(data) if source_shards_by_worker is None else source_shards_by_worker
         )
         self._idx = 0
 
@@ -232,6 +247,8 @@ class CSVIterableDataset(IterableDataset):
 
     """
 
+    _shards_by_worker = True
+
     def __init__(
         self,
         src: str | Sequence[str] | Iterable | Sequence[Iterable],

tests/data/test_shuffle_buffer.py

@@ -90,18 +90,24 @@ def test_worker_sharded_source_is_not_sharded_twice(self):
             self.assertEqual(len(outputs), 40)
             self.assertEqual(set(outputs), set(range(40, 80)))
 
-    def test_explicit_unsharded_source_keeps_outer_worker_partition(self):
-        """Verify explicit unsharded mode preserves outer worker partitioning."""
-        outputs = []
-        for worker_id in range(2):
-            source = _UnshardedMonaiIterable(range(40))
-            buffer = ShuffleBuffer(source, buffer_size=8, seed=7, source_shards_by_worker=False)
-            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)
+    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)))
+            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)