Skip to content

Fix DataLoader with buffer=true, parallel=true, collate=true (#216) - #228

Merged
CarloLucibello merged 1 commit into
mainfrom
cl/fix-buffered-parallel-collate
Jun 14, 2026
Merged

Fix DataLoader with buffer=true, parallel=true, collate=true (#216)#228
CarloLucibello merged 1 commit into
mainfrom
cl/fix-buffered-parallel-collate

Conversation

@CarloLucibello

Copy link
Copy Markdown
Member

Fixes #216.

Problem

A DataLoader configured with buffer=true, parallel=true, and collate=true errored. The original report saw UndefVarError: \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:

julia> DataLoader(rand(Float32, 4, 20); batchsize=5, buffer=true, collate=true, parallel=true) |> first
ERROR: MethodError: no method matching Vector{Vector{Float32}}(::Matrix{Float32})

The parallel buffered path routes results through a RingBuffer whose results channel was typed by the buffer type. With collate=true, getobs! mutates the per-observation buffers (a Vector of per-obs arrays) but returns a freshly batched 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 the RingBuffer.

Fix

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. Recycling stays deferred to the next take!, so results that alias their buffer (collate=false/nothing) remain valid until the consumer asks for the next one — unchanged behavior for those paths.
  • R is read from eltype(BatchView) (its first type parameter, already computed by BatchView), so there is no eager getobs/IO and the channels stay concretely typed. take! is type-stable (@inferred) for concrete data; for tuple-of-array observations R matches eltype(BatchView) exactly, i.e. the same type the serial path already yields.
  • A single-argument RingBuffer(bufs) constructor (defaulting R = B) preserves the existing in-place behavior and the standalone RingBuffer tests.

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.

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>
@CarloLucibello
CarloLucibello merged commit d4232d2 into main Jun 14, 2026
6 of 7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Dataloader buffer=true breaks code

1 participant