Skip to content

Fix stale batch indices in prediction writers for dataloaders without batch indices - #22065

Open
Thanatos9404 wants to merge 2 commits into
Lightning-AI:masterfrom
Thanatos9404:fix/prediction-writer-stale-batch-indices
Open

Fix stale batch indices in prediction writers for dataloaders without batch indices#22065
Thanatos9404 wants to merge 2 commits into
Lightning-AI:masterfrom
Thanatos9404:fix/prediction-writer-stale-batch-indices

Conversation

@Thanatos9404

@Thanatos9404 Thanatos9404 commented Oct 27, 2026

Copy link
Copy Markdown

What does this PR do?

Fixes #22059

When a dataloader's batch indices cannot be inferred (for example an IterableDataset), _store_data_for_prediction_writer returned early without touching current_batch_indices, so BasePredictionWriter.write_on_batch_end received the indices of the previous map-style batch. This happened both for a later dataloader in the same predict() call and across separate predict() calls on the same Trainer.

This resets current_batch_indices to an empty list in that branch, as suggested in the issue, so those batches get [], the same as on a fresh Trainer.

The new test covers both cases. It fails on master with [[0, 1], [2, 3], [2, 3], [2, 3]] and passes with the change. I also ran the prediction writer, prediction loop, fetcher, loop state dict and trainer predict tests locally. CHANGELOG updated.

Thanks @AryanRG for the clear diagnosis; happy to close this if you'd rather send your own fix.

Before submitting
  • Was this discussed/agreed via a GitHub issue? (not for typos and docs)
  • Did you read the contributor guideline, Pull Request section?
  • Did you make sure your PR does only one thing, instead of bundling different changes together?
  • Did you make sure to update the documentation with your changes? (if necessary)
  • Did you write any new necessary tests? (not for typos and docs)
  • Did you verify new and existing tests pass locally with your changes?
  • Did you list all the breaking changes introduced by this pull request?
  • Did you update the CHANGELOG? (not for typos, docs, test updates, or minor internal changes/refactors)

PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:

Reviewer checklist
  • Is this pull request ready for review? (if not, please submit in draft mode)
  • Check that all items from Before submitting are resolved
  • Make sure the title is self-explanatory and the description concisely explains the PR
  • Add labels and milestones (and optionally projects) to the PR so it can be classified

@deependujha
deependujha requested a balanced review from Copilot October 29, 2026 18:38
Copilot AI reviewed Oct 29, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

deependujha
deependujha approved these changes Oct 29, 2026
View reviewed changes

@deependujha deependujha left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, thanks for contributing :)

@codecov-commenter

codecov-commenter commented Oct 29, 2026 •
edited
Loading

Copy link
Copy Markdown

⚠️ Please install the 'codecov app svg image' to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79%. Comparing base (84df182) to head (cb658e6).
✅ All tests successful. No failed tests found.
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

❗ There is a different number of reports uploaded between BASE (84df182) and HEAD (cb658e6). Click for more details.

HEAD has 15673 uploads less than BASE
Flag BASE (84df182) HEAD (cb658e6)
cpu 3562 51
lightning_fabric 1658 0
pytest 1882 0
python3.12 1481 21
python3.12.7 210 3
python3.13 1445 21
python 210 3
lightning 425 6
python3.10 216 3
pytorch_lightning 1479 45
pytorch2.6 295 9
pytest-full 1680 51
pytorch2.12 201 6
pytorch2.7 101 3
pytorch2.9 196 6
pytorch2.8 192 6
pytorch2.13 309 9
pytorch2.10 196 6
pytorch2.11 190 6
Additional details and impacted files 
@@            Coverage Diff            @@
##           master   #22065     +/-   ##
=========================================
- Coverage      87%      79%     -8%     
=========================================
  Files         270      267      -3     
  Lines       24015    23957     -58     
=========================================
- Hits        20855    18826   -2029     
- Misses       3160     5131   +1971

@anishpentyala

anishpentyala commented Oct 31, 2026 •
edited
Loading

Copy link
Copy Markdown

Additional CPU coverage on cb658e6, beyond the existing approval: an independent Trainer harness checked all three writer intervals (batch, epoch, batch_and_epoch), workers 0/1, map-then-iterable, iterable-then-map, map-iterable-map, and separate predict calls on the same Trainer.

All 24 combinations pass on this head. On base 84df182, 12 combinations fail with stale map indices reaching iterable batch callbacks. The other 12 pass, including epoch-only controls. Epoch index lists and the later map loader's real indices remain correct on head.

The PR's prediction-writer tests plus the adjacent prediction-loop and fetcher suites give 112 passed on head. The same test files on base give 111 passed and one inspected stale-indices assertion failure, not a setup error.

Python 3.11, single-process CPU only. No GPU or distributed claim. This verification and comment were AI-assisted

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

Copilot code review Copilot
Copilot review effort
Applies to this pull request for everyone.Learn more about Copilot code review.
Copilot left review comments
@deependujha deependujha deependujha approved these changes
@tchaton tchaton Awaiting requested review from tchaton tchaton is a code owner
@justusschock justusschock Awaiting requested review from justusschock justusschock is a code owner
@ethanwharris ethanwharris Awaiting requested review from ethanwharris ethanwharris is a code owner

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Prediction writers receive stale batch indices for IterableDataset

5 participants

@Thanatos9404 @codecov-commenter @anishpentyala @deependujha