Fix double sharding in ShuffleBuffer

[yinkev]

Closes #7986.

Description

ShuffleBuffer no longer applies worker partitioning twice when its source is already a PyTorch/MONAI IterableDataset.

A MONAI IterableDataset shards its source in its own iterator. ShuffleBuffer then wrapped that already-sharded generator in another IterableDataset, which sharded it again. With two workers this reduced 40 unique source items to 20 outputs. Ordinary iterable sources retain the existing shuffle-then-partition path, and transforms remain exactly-once.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

Validation

• deterministic regression fails on upstream/dev due to double sharding;
• python -m pytest tests/data/test_shuffle_buffer.py tests/data/test_iterable_dataset.py -q: 5 passed;
• real two-worker reproduction: 40 inputs, 40 unique outputs;
• ruff check on both touched files;
• git diff --check.

[coderabbitai]

📝 Walkthrough

Walkthrough

ShuffleBuffer.__iter__ preserves the existing worker partitioning of MONAI IterableDataset sources by iterating generated items directly and applying transforms inline. Other sources continue through the wrapper-based path. A regression test mocks multi-worker information and verifies that wrapping an IterableDataset does not shard items twice.

Estimated code review effort: 2 (Simple) | ~10 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The changes address #7986 by avoiding a second partitioning step for IterableDataset sources, which preserves all patches.
Out of Scope Changes check	✅ Passed	The diff stays focused on the ShuffleBuffer sharding fix and its regression test with no unrelated changes.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Title check	✅ Passed	The title is concise and accurately summarizes the main change: preventing double sharding in ShuffleBuffer.
Description check	✅ Passed	The description covers the required summary and types of changes sections and adds useful validation details.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (1)

monai/data/iterable_dataset.py (1)

127-128: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add Google-style docstrings to the changed definitions.

• monai/data/iterable_dataset.py#L127-L128: document the yielded item with a Yields: section.
• tests/data/test_shuffle_buffer.py#L42-L42: add a concise docstring describing the regression scenario and assertions.

As per path instructions, every definition must have an appropriate Google-style docstring.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@monai/data/iterable_dataset.py` around lines 127 - 128, Document the changed
definition in monai/data/iterable_dataset.py at lines 127-128 with a
Google-style docstring that includes a Yields: section describing the yielded
item. Also add a concise Google-style docstring to the regression test
definition in tests/data/test_shuffle_buffer.py at line 42 describing the
scenario and assertions.

Source: Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@monai/data/iterable_dataset.py`:
- Around line 127-128: Document the changed definition in
monai/data/iterable_dataset.py at lines 127-128 with a Google-style docstring
that includes a Yields: section describing the yielded item. Also add a concise
Google-style docstring to the regression test definition in
tests/data/test_shuffle_buffer.py at line 42 describing the scenario and
assertions.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 2540d5b4-e2cd-4ba1-b2d5-c7952b78608e

📥 Commits

Reviewing files that changed from the base of the PR and between 3ee058b and 52ce061.

📒 Files selected for processing (2)

• monai/data/iterable_dataset.py
• tests/data/test_shuffle_buffer.py

[yinkev]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[yinkev]

Replaced by #9020 to conform to this repository's documented contributor workflow and branch-naming convention. The replacement uses the same signed commit and identical code diff: #9020