Fix double sharding in ShuffleBuffer by yinkev · Pull Request #9020 · Project-MONAI/MONAI

yinkev · 2026-07-25T19:45:51Z

Description

ShuffleBuffer now distinguishes sources that already partition their stream using torch.utils.data.get_worker_info. MONAI IterableDataset sources are recognized automatically, while an explicit source_shards_by_worker override supports generic PyTorch or unusual custom iterable sources.

The default auto-detects MONAI iterable sources, so the GridPatchDataset pipeline reported in #7986 works without a caller change. True prevents a second partition for any worker-aware source; False preserves the outer worker partition for an explicitly unsharded iterable source.

This avoids using class identity as a proxy for sharding behavior: MONAI subclasses can be unsharded, while generic PyTorch iterable sources can shard themselves.

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

the original double-sharding reproduction loses samples when an already-sharded source is partitioned again;
tests cover both explicitly worker-sharded MONAI/PyTorch sources and an unsharded MONAI subclass whose default behavior must remain unchanged;
affected data tests: 7 passed;
real two-worker fork-context DataLoader contract: the default MONAI-source path returns all 40 inputs exactly once;
Black, isort, Ruff, DCO, and git diff --check pass.

Signed-off-by: kyinhub <kevinpyin@gmail.com>

coderabbitai · 2026-07-25T19:46:10Z

📝 Walkthrough

Walkthrough

ShuffleBuffer now detects or accepts whether its source already shards data across workers, choosing direct iteration for sharded sources and wrapper-based partitioning otherwise. Tests cover default detection, explicit worker-sharded sources, and unsharded sources.

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

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The code and tests address `#7986` by preventing ShuffleBuffer from double-sharding worker-aware iterable sources.
Out of Scope Changes check	✅ Passed	The changes stay focused on ShuffleBuffer sharding behavior, docstrings, and regression tests with no obvious unrelated additions.
Title check	✅ Passed	The title clearly summarizes the main change: fixing double sharding in ShuffleBuffer.
Description check	✅ Passed	The description matches the template with an issue link, a change summary, types of changes, and added 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

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

coderabbitai

🧹 Nitpick comments (1)

monai/data/iterable_dataset.py (1)
125-132: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Document transform failures.

Add a Google-style Raises section for RuntimeError from apply_transform.

As per path instructions, “Docstrings should be present for all definition which describe each variable, return value, and raised exception in the appropriate section of the Google-style of docstrings.”
🤖 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 125 - 132, Add a Google-style
Raises section to the docstring of the method yielding shuffled items,
documenting that apply_transform can raise RuntimeError. Keep the existing
description and Yields documentation unchanged.
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 125-132: Add a Google-style Raises section to the docstring of the
method yielding shuffled items, documenting that apply_transform can raise
RuntimeError. Keep the existing description and Yields documentation unchanged.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: f032a9a6-a682-4939-bf36-d6b81c16e970

📥 Commits

Reviewing files that changed from the base of the PR and between 3ee058b and 4000d9a.

📒 Files selected for processing (2)

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

yinkev · 2026-07-25T23:24:34Z

@coderabbitai review

coderabbitai · 2026-07-25T23:24:38Z

✅ 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.

coderabbitai

Actionable comments posted: 1

🧹 Nitpick comments (1)

monai/data/iterable_dataset.py (1)
105-120: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add Google-style method docstrings.

monai/data/iterable_dataset.py#L105-L120: document constructor arguments and initialization behavior.

tests/data/test_shuffle_buffer.py#L30-L44: document yielded values for both helper iterators.

As per path instructions, “Docstrings should be present for all definition[s] … in the appropriate section of the Google-style of docstrings.”
🤖 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 105 - 120, Document the
IterableDataset constructor with a Google-style docstring covering each argument
and its initialization behavior. In tests/data/test_shuffle_buffer.py lines
30-44, add Google-style docstrings to both helper iterator definitions
describing the values they yield; make the requested documentation changes at
both sites.
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.

Inline comments:
In `@tests/data/test_shuffle_buffer.py`:
- Around line 93-104: Update the test description and method name in
test_unsharded_source_keeps_outer_worker_partition to describe the explicitly
configured source_shards_by_worker=False behavior rather than claiming it
verifies the default. Keep the test setup and assertions unchanged.

---

Nitpick comments:
In `@monai/data/iterable_dataset.py`:
- Around line 105-120: Document the IterableDataset constructor with a
Google-style docstring covering each argument and its initialization behavior.
In tests/data/test_shuffle_buffer.py lines 30-44, add Google-style docstrings to
both helper iterator definitions describing the values they yield; make the
requested documentation changes at both sites.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 4122d582-cc91-43b8-ae0f-fb344ed8ff38

📥 Commits

Reviewing files that changed from the base of the PR and between 4000d9a and bfbaacd.

📒 Files selected for processing (2)

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

Signed-off-by: kyinhub <kevinpyin@gmail.com>

coderabbitai

Actionable comments posted: 1

🤖 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.

Inline comments:
In `@monai/data/iterable_dataset.py`:
- Around line 105-120: Add a Google-style docstring directly inside
ShuffleBuffer.__init__ documenting all constructor arguments, including
source_shards_by_worker, and the constructor’s return value. Keep the existing
initialization behavior unchanged.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 2bcf7957-7fe6-4929-a20e-95b51dd7f4af

📥 Commits

Reviewing files that changed from the base of the PR and between bfbaacd and e525693.

📒 Files selected for processing (2)

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

🚧 Files skipped from review as they are similar to previous changes (1)

tests/data/test_shuffle_buffer.py

Signed-off-by: kyinhub <kevinpyin@gmail.com>

fix(data): avoid double sharding shuffle buffer

4000d9a

Signed-off-by: kyinhub <kevinpyin@gmail.com>

yinkev requested review from KumoLiu, Nic-Ma and ericspod as code owners July 25, 2026 19:45

yinkev mentioned this pull request Jul 25, 2026

Fix double sharding in ShuffleBuffer #9014

Closed

7 tasks

coderabbitai Bot reviewed Jul 25, 2026

View reviewed changes

yinkev marked this pull request as draft July 25, 2026 20:59

coderabbitai Bot reviewed Jul 25, 2026

View reviewed changes

Comment thread tests/data/test_shuffle_buffer.py Outdated

fix(data): make ShuffleBuffer sharding explicit

e525693

Signed-off-by: kyinhub <kevinpyin@gmail.com>

yinkev force-pushed the 7986-shuffle-buffer-double-shard branch from bfbaacd to e525693 Compare July 26, 2026 02:21

yinkev marked this pull request as ready for review July 26, 2026 06:00

coderabbitai Bot reviewed Jul 26, 2026

View reviewed changes

Comment thread monai/data/iterable_dataset.py

docs(data): document ShuffleBuffer constructor

cb3a713

Signed-off-by: kyinhub <kevinpyin@gmail.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Fix double sharding in ShuffleBuffer#9020

Fix double sharding in ShuffleBuffer#9020
yinkev wants to merge 3 commits into
Project-MONAI:devfrom
yinkev:7986-shuffle-buffer-double-shard

yinkev commented Jul 25, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Jul 25, 2026 •

edited

Loading

Walkthrough

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Uh oh!

yinkev commented Jul 25, 2026

Uh oh!

coderabbitai Bot commented Jul 25, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

yinkev commented Jul 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Types of changes

Validation

Uh oh!

coderabbitai Bot commented Jul 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

yinkev commented Jul 25, 2026

Uh oh!

coderabbitai Bot commented Jul 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

yinkev commented Jul 25, 2026 •

edited

Loading

coderabbitai Bot commented Jul 25, 2026 •

edited

Loading

coderabbitai Bot commented Jul 25, 2026 •

edited

Loading