From 35b092bcbb726572fd0b573d5d09902d28f03f1e Mon Sep 17 00:00:00 2001 From: Parth Vadera Date: Fri, 23 Oct 2026 16:32:51 +0100 Subject: [PATCH 1/2] Use deque for O(1) batch queue pops in _PrefetchDataFetcher Signed-off-by: Parth Vadera --- src/lightning/pytorch/loops/fetchers.py | 7 ++++--- tests/tests_pytorch/loops/test_fetchers.py | 3 ++- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/lightning/pytorch/loops/fetchers.py b/src/lightning/pytorch/loops/fetchers.py index 92ec95a9e2f58..5ba5be03bd021 100644 --- a/src/lightning/pytorch/loops/fetchers.py +++ b/src/lightning/pytorch/loops/fetchers.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from collections import deque from collections.abc import Iterator from typing import Any, Optional @@ -98,7 +99,7 @@ def __init__(self, prefetch_batches: int = 1) -> None: if prefetch_batches < 0: raise ValueError("`prefetch_batches` should at least be 0.") self.prefetch_batches = prefetch_batches - self.batches: list[Any] = [] + self.batches: deque[Any] = deque() @override def __iter__(self) -> "_PrefetchDataFetcher": @@ -122,7 +123,7 @@ def __next__(self) -> _ITERATOR_RETURN: if self.batches: # there are pre-fetched batches already from a previous `prefetching` call. # consume one - batch = self.batches.pop(0) + batch = self.batches.popleft() try: # refill the consumed batch self.batches.append(super().__next__()) @@ -140,7 +141,7 @@ def __next__(self) -> _ITERATOR_RETURN: @override def reset(self) -> None: super().reset() - self.batches = [] + self.batches.clear() class _DataLoaderIterDataFetcher(_DataFetcher): diff --git a/tests/tests_pytorch/loops/test_fetchers.py b/tests/tests_pytorch/loops/test_fetchers.py index 04f6fc23aa8df..95d6174bc73e5 100644 --- a/tests/tests_pytorch/loops/test_fetchers.py +++ b/tests/tests_pytorch/loops/test_fetchers.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from collections import Counter +from collections import Counter, deque from collections.abc import Iterator from typing import Any @@ -51,6 +51,7 @@ def __getitem__(self, idx): def test_prefetch_iterator(multiple_iterables, dataset_cls, prefetch_batches): fetcher = _PrefetchDataFetcher(prefetch_batches=prefetch_batches) assert fetcher.prefetch_batches == prefetch_batches + assert isinstance(fetcher.batches, deque) if multiple_iterables: loader = CombinedLoader([DataLoader(dataset_cls()), DataLoader(dataset_cls())]) From 60ff6ed3e5366018d9282bf5a80f4e8f449ab9d7 Mon Sep 17 00:00:00 2001 From: Parth Vadera Date: Thu, 29 Oct 2026 21:36:26 +0000 Subject: [PATCH 2/2] test: assert reset() empties pending prefetched batches instead of isinstance Suggested-by: Deependu Jha Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Parth Vadera --- tests/tests_pytorch/loops/test_fetchers.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/tests_pytorch/loops/test_fetchers.py b/tests/tests_pytorch/loops/test_fetchers.py index 95d6174bc73e5..2e57cbf8801ba 100644 --- a/tests/tests_pytorch/loops/test_fetchers.py +++ b/tests/tests_pytorch/loops/test_fetchers.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from collections import Counter, deque +from collections import Counter from collections.abc import Iterator from typing import Any @@ -51,7 +51,6 @@ def __getitem__(self, idx): def test_prefetch_iterator(multiple_iterables, dataset_cls, prefetch_batches): fetcher = _PrefetchDataFetcher(prefetch_batches=prefetch_batches) assert fetcher.prefetch_batches == prefetch_batches - assert isinstance(fetcher.batches, deque) if multiple_iterables: loader = CombinedLoader([DataLoader(dataset_cls()), DataLoader(dataset_cls())]) @@ -78,6 +77,10 @@ def generate(): # validate reset works properly. assert generate() == expected assert fetcher.fetched == 3 + iter(fetcher) + assert len(fetcher.batches) == (0 if dataset_cls is SizedDataset else min(prefetch_batches, 3)) + fetcher.reset() + assert not fetcher.batches @pytest.mark.parametrize("multiple_iterables", [False, True])