Fix DataLoader with buffer=true, parallel=true, collate=true (#216) - #228
Merged
Conversation
The parallel buffered DataLoader path routed results through a `RingBuffer`
whose results channel was typed by the *buffer* type. With `collate=true`,
`getobs!` mutates the per-observation buffers but returns a freshly batched
array of a different type, so putting it into the channel errored
(MethodError / convert error).
Parametrize `RingBuffer{B,R}` so the buffer type `B` and result type `R` are
tracked separately: the results channel now carries `(buffer, result)` pairs
and recycles the buffer (not the result) back into the pool, keeping the
aliasing safety of the `collate=false`/`nothing` paths intact. `R` is read
from `eltype(BatchView)` (no eager `getobs`), so the channel stays concretely
typed and `take!` remains type-stable for concrete data.
Add a regression test covering the array, tuple, and custom-collate cases.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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
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 #216.
Problem
A
DataLoaderconfigured withbuffer=true,parallel=true, andcollate=trueerrored. The original report sawUndefVarError: \A` not defined in `MLUtils``; that came from the FLoops-based loader and disappeared with #215, but the same configuration still failed afterwards with a type error:The parallel buffered path routes results through a
RingBufferwhose results channel was typed by the buffer type. Withcollate=true,getobs!mutates the per-observation buffers (aVectorof per-obs arrays) but returns a freshlybatched array of a different type. Putting that collated result into a channel typed for the buffer failed. The serial path was unaffected because it does not go through theRingBuffer.Fix
Parametrize
RingBuffer{B,R}so the buffer typeBand result typeRare tracked separately:(buffer, result)pairs and recycles the buffer (not the result) back into the pool. Recycling stays deferred to the nexttake!, so results that alias their buffer (collate=false/nothing) remain valid until the consumer asks for the next one — unchanged behavior for those paths.Ris read fromeltype(BatchView)(its first type parameter, already computed byBatchView), so there is no eagergetobs/IO and the channels stay concretely typed.take!is type-stable (@inferred) for concrete data; for tuple-of-array observationsRmatcheseltype(BatchView)exactly, i.e. the same type the serial path already yields.RingBuffer(bufs)constructor (defaultingR = B) preserves the existing in-place behavior and the standaloneRingBuffertests.Tests
Added a regression test ("buffer + collate + parallel issue 216") covering plain-array (with serial-vs-parallel correctness), and a custom collate function over a non-array container. Full test suite passes locally.