Release train DataLoader workers before validation#890
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bad13082d3
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
bad1308 to
58ee117
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 58ee11752d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
373a8a7 to
053c886
Compare
Guarantee validation DataLoader worker release via try/finally Wrap the validation loop in try/finally and release the val DataLoader iterator (val_dataloader_iter = None) in the finally block, replacing the bare `del` at the end of the loop. Functionally equivalent to the previous `del` (both drop the sole reference so CPython ref-counting runs _MultiProcessingDataLoaderIter._shutdown_workers immediately), but the release is now structurally guaranteed: it still runs if the validation loop exits early via an exception or a future break/return, and it reads symmetrically with the train-side infinite_train_dataloader.reset(). Verified on a 2-GPU / num_workers=8 SLURM run over a small COCO instance-seg dataset: concurrent DataLoader workers stay at 16 (2 ranks x 8) through every validation phase and never double to 32 (the pre-fix behavior), matching the original fix. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> fix codex reviews Address review: robust val iterator release and zero-batch guard - Release the validation iterator with `val_dataloader_iter = None` instead of `del` in the finally block: functionally identical, matches the PR description, and won't raise NameError if the iter() assignment is ever moved inside the try. - Guard the train_dataloader_position modulo with max(train_num_batches, 1) to avoid ZeroDivisionError on a zero-length train dataloader. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Revert val iterator release to del to satisfy mypy The `val_dataloader_iter = None` variant violates the variable's inferred iterator type under mypy. Use `del val_dataloader_iter` in the finally block instead (functionally identical). Keeps the max(train_num_batches, 1) zero-batch guard. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Revert val iterator release to del to satisfy mypy The `val_dataloader_iter = None` variant violates the variable's inferred iterator type under mypy. Use `del val_dataloader_iter` in the finally block instead (functionally identical). Keeps the max(train_num_batches, 1) zero-batch guard. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
053c886 to
5dc298a
Compare
|
/review |
6332734 to
117a7af
Compare
There was a problem hiding this comment.
Tbh, I don't think we can merge it like this. I don't have an idea at the moment either on how to handle this, but if I have a calm minute in the next few days I will try to figure something out. The sampler for the train dataloader needs to somehow be resumable.
| # the current epoch's remaining batches, but we make no | ||
| # per-epoch coverage guarantees anyway. The next next() call | ||
| # re-creates the iterator and respawns the workers. | ||
| infinite_train_dataloader.reset() |
There was a problem hiding this comment.
Now there is basically no guarantee that if you set steps=len(dataset) // batch_size that every sample gets seen once. I am not sure if this is desirable tbh and would prefer a different solution, since most people would expect that to be the case. Also we still log the "epoch".
|
Note: we would go with improving the batch sampler later on and leave this PR as is atm. (Issue backlogged with a plan on Linear) |
2 participants
What has changed and why?
During task training, two DataLoader worker pools could be alive at the same time throughout every validation phase.
InfiniteCycleIteratorcaches its underlying iterator and keeps it alive across training steps, so the train workers did not shut down before validation. When validation started,iter(val_dataloader)spawned a second pool while the train pool was still held. Because the validation loader uses the samenum_workersas training, the number of live worker processes doubled during validation (for example, 2 GPUs x 8 workers = 16 -> 32).This PR ensures only one task DataLoader worker pool is alive at a time:
num_workers > 0,infinite_train_dataloader.reset()drops the cached train iterator before validation starts, allowing its workers to shut down before validation workers spawn.try/finallyand deleted in thefinally, so validation workers are released even if validation exits early or raises.persistent_workerspassed throughloader_argsis now ignored for both task train and validation dataloaders. Persistent workers are owned by the DataLoader and survive iterator shutdown, so allowing this override would defeat the release-before-validation behavior.How has it been tested?
Ran a 2-GPU /
num_workers=4job under SLURM on a small COCO instance-segmentation subset (ltdetrv2-seg-s, 30 steps, validating at steps 10/20/30 plus the final step). Apsutilmonitor sampled the number of live DataLoader worker processes, matched by theirpt_data_workerprocess name and isolated to the job by run name, every 0.25 s.Did you update CHANGELOG.md?
Did you update the documentation?
🤖 Generated with Claude Code