Fix double iteration bug when resumed from a checkpoint.

[sudiptob2]

What does this PR do?

This PR fixes the double iter() bug discussed in #19427

Fixes #19427

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

---

📚 Documentation preview 📚: https://pytorch-lightning--20775.org.readthedocs.build/en/20775/

[sudiptob2]

@Borda Thanks for the review! Let me know if there's anything you'd like me to change. Otherwise, can we go ahead and merge this?

[sudiptob2]

Hey, @Borda just wondering if anything is blocking this PR from merging?

[sudiptob2]

@Borda, I noticed a few mypy checks are failing, but they don't seem related to the changes in this PR. Should we ignore them for now? Let me know the next steps.

[sudiptob2]

@bhimrazy, thanks for approving the PR. I can see that a few checks are timed out. Any idea why?

[sudiptob2]

Just checking in, any hope of moving this forward?

[deependujha]

Hi @sudiptob2

Sorry for the delay and thanks for your patience.
I’m checking out the PR now and will update you shortly.

[Copilot]

Pull Request Overview

This PR fixes a double iteration bug that occurs when resuming training from a checkpoint with IterableDatasets. The issue was that iter() was being called twice - once during setup and again during epoch start - causing data to be skipped when resuming from checkpoints.

• Added a new is_resuming property to track when training is being resumed from a checkpoint
• Modified the iteration logic in TrainingEpochLoop to avoid double iteration when resuming
• Added comprehensive test coverage for the checkpoint resumption scenario

Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
src/lightning/pytorch/loops/loop.py	Added is_resuming property and state tracking for checkpoint resumption
src/lightning/pytorch/loops/training_epoch_loop.py	Modified iteration logic to check resumption state before calling iter()
tests/tests_pytorch/loops/test_double_iter_in_iterable_dataset.py	Added test case to verify checkpoint resumption works correctly with IterableDatasets
.github/workflows/ci-tests-pytorch.yml	Increased timeout to accommodate additional test execution time

Comments suppressed due to low confidence (2)

tests/tests_pytorch/loops/test_double_iter_in_iterable_dataset.py:66

• The function name test_resume_training_with is incomplete and unclear. It should be more descriptive, such as test_resume_training_with_iterable_dataset or test_resume_training_with_checkpoint.

def test_resume_training_with(tmp_path):

tests/tests_pytorch/loops/test_double_iter_in_iterable_dataset.py:69

• The variable name max_epoch should be max_epochs to be consistent with the parameter name used in the train_model function and PyTorch Lightning conventions.

    max_epoch = 2

[deependujha]

Hi @sudiptob2,

I’m seeing an odd issue: all tests pass on macOS 14 with Python 3.12 and 2.7 in under 30 minutes, but the job still hangs until the 50‑minute timeout and then fails.

I suspect the issue to be in the test file, possibly related toqueue handling or multiprocessing (spawn/fork).

[deependujha]

Thanks, @sudiptob2, for your contribution and for making it even better!

[sudiptob2]

@deependujha Thanks to you for the hard work with testing and CI fixes. Really appreciate 🚀