Fix thread leak in threaded_generator when consumers exit early#721
Open
discobot wants to merge 1 commit into
Open
Fix thread leak in threaded_generator when consumers exit early#721discobot wants to merge 1 commit into
discobot wants to merge 1 commit into
Conversation
The fill_queue thread blocks forever on q.put once its bounded queue fills up if the consuming generator is abandoned before exhaustion, which happens every time a data loader epoch iterator isn't fully consumed. This leaks one thread per abandoned generator and CPU thread count grows steadily over a training run. The producer now checks a stop event while putting, and the consumer sets it and drains the queue on exit, so the thread shuts down as soon as the consumer goes away.
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #703.
The issue correctly pinpoints
threaded_generator. The mechanism: itsfill_queuedaemon thread feeds a bounded queue (maxsize=16by default) with a plain blockingq.putand has no stop signal, so if the consumer drops the wrapper generator before exhausting it, the thread blocks onputforever once the queue fills.What makes this grow steadily during training: the data loaders (
data_loader.pyandcomposable/data_loader.py) createnum_threadsof these wrappers per batch-iteration and abandon them whenever an epoch iterator isn't fully consumed — stopping mid-epoch, hitting max duration, recreating the iterator — stranding another batch of threads each time.This adds a stop event: the producer puts with a short timeout and re-checks the event, and the consumer sets it and drains the queue from a
finallyblock, which also runs onGeneratorExitwhen the wrapper is closed or garbage collected. Producer threads now exit promptly once their consumer goes away. Ordering, sentinel termination, and exception forwarding viaOLMoThreadErrorare unchanged.Testing: new regression test in
src/test/utils_test.py— abandoning wrappers mid-iteration leaks all 8fill_queuethreads on main and none with this change — plus coverage for in-order full consumption and exception forwarding.src/test/data/data_loader_test.pyandsrc/test/data/composable/data_loader_test.pypass.