Skip to content

Use deque for batch queue in _PrefetchDataFetcher - #22052

Open
deep1blue wants to merge 2 commits into
Lightning-AI:masterfrom
deep1blue:perf/prefetch-data-fetcher-deque
Open

deep1blue wants to merge 2 commits into
Lightning-AI:masterfrom
deep1blue:perf/prefetch-data-fetcher-deque

Conversation

@deep1blue

@deep1blue deep1blue commented Oct 23, 2026 •

Copy link
Copy Markdown

What does this PR do?

_PrefetchDataFetcher in src/lightning/pytorch/loops/fetchers.py uses self.batches as a FIFO queue for prefetched batches on unsized iterables (appending at the tail and popping from the head on each __next__() call).

This cleanup switches self.batches from a list (pop(0) / = []) to collections.deque (popleft() / clear()) to use the idiomatic FIFO container.

Verification

  • pytest tests/tests_pytorch/loops/test_fetchers.py (98 passed)

Signed-off-by: Parth Vadera <parthvadera@gmail.com>
@deependujha

Copy link
Copy Markdown
Collaborator

Thanks, the change is correct and deque is the more idiomatic structure for a FIFO here, so happy to take it as a cleanup.

Two notes:

  • The perf framing overstates it: prefetch_batches defaults to 1 and this path only runs for unsized iterables (sized loaders skip prefetching entirely). With B this small, list.pop(0) vs deque.popleft() is a few ns per step. Could you reword the description as a readability cleanup rather than a perf fix?

  • Could you drop assert isinstance(fetcher.batches, deque)? It pins an implementation detail. If you want a test, asserting that reset() empties pending prefetched batches would cover the one behavioural line that changed.

@codecov-commenter

codecov-commenter commented Oct 29, 2026 •

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 (60ff6ed).
✅ 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 (60ff6ed). Click for more details.

HEAD has 16045 uploads less than BASE
Flag BASE (84df182) HEAD (60ff6ed)
cpu 3646 51
lightning_fabric 1701 0
pytest 1930 0
python3.12 1516 21
python3.12.7 215 3
python3.13 1481 21
python 213 3
lightning 435 6
python3.10 221 3
pytorch_lightning 1510 45
pytorch2.6 299 9
pytest-full 1716 51
pytorch2.12 205 6
pytorch2.7 103 3
pytorch2.9 199 6
pytorch2.8 196 6
pytorch2.13 317 9
pytorch2.10 201 6
pytorch2.11 196 6
Additional details and impacted files
@@            Coverage Diff            @@
##           master   #22052     +/-   ##
=========================================
- Coverage      87%      79%     -8%     
=========================================
  Files         270      267      -3     
  Lines       24015    23957     -58     
=========================================
- Hits        20855    18826   -2029     
- Misses       3160     5131   +1971     

…instance

Suggested-by: Deependu Jha <deependujha@users.noreply.github.com>

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Parth Vadera <parthvadera@gmail.com>
@deep1blue deep1blue changed the title Use deque for O(1) batch queue pops in _PrefetchDataFetcher Use deque for batch queue in _PrefetchDataFetcher Oct 29, 2026
@deep1blue

Copy link
Copy Markdown
Author

Thanks @deependujha! Updated in 60ff6ed:

  • Reworded the PR title and description as a readability cleanup using the idiomatic FIFO container.
  • Dropped assert isinstance(fetcher.batches, deque) and replaced it with an assertion in test_prefetch_iterator verifying that fetcher.reset() empties pending prefetched batches.

Verified locally with targeted tests (nice pytest tests/tests_pytorch/loops/test_fetchers.py: 98 passed).

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants