Allow parallel=n to set the number of loading workers

[CarloLucibello]

Allow to set the number of parallel workers in the DataLoader (previously other serial behavior or num workers = num threads).

Fixes #210 (DataLoader parallel loading uses too much memory).

What

Extends the parallel kwarg on DataLoader/eachobs from Bool to Union{Bool,Integer}:

parallel=	meaning
false / 0	serial (current default)
true	Threads.nthreads() workers (current behavior)
n::Integer	exactly n worker tasks
1	one background worker overlapping load with the iterating task

This follows PyTorch's convention (num_workers=0 ⇒ main-process/serial), folded onto the existing kwarg rather than adding a separate num_workers — so there's no contradictory state like parallel=false, num_workers=8.

Why

In the parallel path each in-flight worker holds one whole (collated) batch, so peak memory is dominated by nworkers × batchsize, not the prefetch channel depth. The worker count was hardcoded to Threads.nthreads(), so a 16-thread session building ~600 MB batches kept ~16–32 batches live → the 40 GB RSS in #210. An experiment confirmed channelsize is not a useful memory lever (lowering it didn't help, sometimes hurt), while the worker count is — so this exposes the worker count and leaves channelsize derived from it.

How

• _nworkers resolves the kwarg, defined by dispatch (not ==) so Bool <: Integer doesn't conflate true with 1.
• Threaded through eachobsparallel → Loader → basesize = length ÷ nworkers.
• channelsize now defaults to nworkers (was nthreads()), so the prefetch buffer scales with the worker count instead of the global thread count.
• The raw parallel value is stored unchanged for faithful printing / reconstruction (parallel=true doesn't silently round-trip to parallel=8).
• parallel=true resolves to nthreads(), so all existing defaults and behavior are unchanged. Widening Bool → Union{Bool,Integer} is backward compatible.

Benchmark

Reproducing the #210 setup on a 16-core / 32-thread Threadripper (122 GiB RAM): 40 batches × 1024 images (224×224×3 Float32, ~600 MB/collated-batch), collate=true, rand-based getobs (CPU-bound, no I/O). Each cell is a fresh process so Sys.maxrss is clean. (getobs! defined so buffer=true reuses per-observation image buffers.)

Peak RSS (GB):

parallel	buffer	t=4	t=8	t=16
false	false	2.6	2.6	2.4
false	true	1.5	1.5	1.5
1	false	2.5	2.4	2.4
1	true	2.1	3.3	3.3
2	false	4.0	4.0	4.1
2	true	4.4	5.0	4.4
4	false	6.6	6.6	6.6
4	true	9.0	8.5	7.9
auto	false	6.6	9.1	15.3
auto	true	8.5	11.4	23.5

Time in seconds (GC% in parens):

parallel	buffer	t=4	t=8	t=16
false	false	6.98 (25%)	6.51 (21%)	7.04 (22%)
false	true	5.35 (8%)	5.26 (8%)	5.75 (10%)
1	false	7.30 (24%)	7.95 (23%)	7.28 (23%)
1	true	5.41 (8%)	5.41 (7%)	5.46 (8%)
2	false	6.29 (26%)	6.00 (22%)	6.06 (24%)
2	true	5.46 (15%)	5.33 (13%)	5.41 (13%)
4	false	5.55 (29%)	5.51 (21%)	5.78 (22%)
4	true	5.62 (16%)	5.40 (13%)	5.59 (13%)
auto	false	5.58 (29%)	4.82 (23%)	5.67 (21%)
auto	true	5.60 (16%)	6.26 (14%)	5.35 (16%)

Takeaways:

1. parallel=n is the memory dial. With buffer=false, peak RSS tracks nworkers × batchsize: 2.5 (serial) → 2.4 (1) → 4 (2) → 6.6 (4) → 15.3 GB (auto @ 16 threads). parallel=1/2 stay bounded regardless of thread count. auto is what blows up — and worsens with thread count.
2. buffer=true is double-edged and must be paired with low nworkers. Serial / parallel=1: it's the lowest memory (1.5 GB, preallocation reuse, ~ the SimpleLoader in DataLoader uses way too much memory #210). But auto, buffer=true is the highest (23.5 GB), because the buffer pool is (nworkers+1) whole batches of per-observation buffers held for the whole iteration, on top of the in-flight collated outputs.
3. Timings are a narrow, noisy band (4.8–8.0 s, ~15% run-to-run). getobs is pure CPU with no I/O to overlap and only 40 batches, so parallel barely moves wall-clock here — treat RSS as the robust number. The consistent time win is buffer=true (halves GC from ~21–29% to ~7–16%). On real I/O-bound loading, higher parallel would help throughput, making the memory/throughput trade-off the actual decision.
4. Sweet spot for large batches: parallel=2, buffer=true — ~5.4 s, ~13% GC, ~4–5 GB, versus auto's 15–23 GB.

Tests

New parallel worker count testset: _nworkers mapping incl. the true/1 distinction, :serial/:parallel codepath dispatch, raw-value round-trip, and that every observation is loaded exactly once for parallel ∈ (false, 0, 1, 2, true) across unbuffered / buffered / batched paths. Also verified parallel=n composes with the buffer=true + collate=true path from #216/#228 (stays @inferred-stable).

🤖 Generated with Claude Code