Process-parallel DataLoader (num_workers) support for Dataset#64
Merged
Conversation
CarloLucibello
force-pushed
the
cl/dataloader-numworkers-guide-and-fix
branch
from
97da41a to
7f5d946
Compare
CarloLucibello
force-pushed
the
cl/dataloader-numworkers-guide-and-fix
branch
from
7f5d946 to
3e208fd
Compare
`MLUtils.DataLoader(ds::Dataset; num_workers=N)` (MLUtils >= 0.4.10) spreads `getobs` over N worker processes, each with its own CPython interpreter and GIL, so reads scale where thread parallelism cannot (a shared GIL serializes them). The loader serializes its data container from a background feeder task that does not hold the GIL; a `Dataset` serializes by calling `pickle`/PythonCall, which segfaults off the main task. Fix it with a `DistributedDataset` wrapper that precomputes the pickle bytes on the calling (GIL-holding) task and ships those bytes — pure Julia, safe from any thread — reconstructing a live `Dataset` on the worker. The `DataLoader(::Dataset)` method installs the wrapper automatically when `num_workers > 0`; serial/`parallel` loaders are untouched. `Dataset` itself is unchanged. An `@info` notes when an in-memory dataset is first materialized to a temporary Arrow file so it can pickle by reference. Tests: a DistributedDataset round-trip plus an end-to-end num_workers DataLoader over an in-memory dataset; MLUtils is added as a test dependency. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Split the combined MLUtils section into a focused integration section and a standalone "Data loaders" section covering feeding a Dataset to DataLoader, on-the-fly vs. materialized loading, and process-parallel `num_workers` loading (with DistributedDataset). Document DistributedDataset in the API reference. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The GIL-safe, feeder-thread-proof serialization needed for `DataLoader(ds; num_workers=N)` now lives in MLUtils (>= 0.4.11): it serializes the data container on the main task and loads this package on the workers. That makes the interim `DistributedDataset` wrapper and the `DataLoader(::Dataset)` hook redundant here, and additionally makes `mapobs`/`ObsView`-wrapped datasets work under `num_workers`. - Remove src/dataloader.jl (DistributedDataset + the DataLoader hook) - Drop the export and the now-unused `numobs` import - Bump MLUtils compat to 0.4.11 (main + test) - Update guide, API docs, changelog, and tests accordingly The `Dataset` Serialization methods (pickle-by-reference) are unchanged -- they are exactly what MLUtils' main-task serialization invokes. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- Fix the `[Data Loaders](@ref)` cross-reference (header case) that failed the Documentation CI job. - Trim the guide's Data Loaders section to defer loader mechanics to the MLUtils guide; keep only the `Dataset`-specific notes. - Move the Flux MNIST example into examples/flux_mnist/ with its own Project.toml, and switch it to process-parallel loading via `num_workers` (package-free transform; one-hot/loss on the main process). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
DataLoader(...; num_workers=N) spawns Distributed workers that each `using PythonCall`. Without help, every worker re-resolves the same CondaPkg env in lockstep: they serialize on CondaPkg's file lock (noisy "Waiting for lock to be freed") and redo work the parent already did. Exporting the already-resolved interpreter in `__init__` lets the workers (which inherit ENV) use it directly and skip CondaPkg entirely. Mirrors the CI-only hack inside PythonCall; guarded on `CTX.which === :CondaPkg` and with `get!` so a user-set interpreter is left untouched. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- Rename examples/flux_mnist/main.jl -> flux_mnist.jl and add a warm-up epoch, so the reported per-config timings exclude Julia's JIT compilation. - Add PyTorch counterparts (a 1:1 port and an idiomatic HF version) and a README with the Julia vs PyTorch timing tables and takeaways (Apple M1 Pro, CPU, 4 epochs). Headline: materializing into memory is the big win; num_workers/process parallelism does not pay off for this toy MLP. - Point the example links in docs/src/guide.md at the new path. - gitignore __pycache__; add DataStructures compat bound. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Split the Data Loaders section into Iterating batches / Materializing / num_workers, with two verified jldoctests (batch structure and the mapobs -> (input, target) pattern) plus julia-repl Hub and num_workers examples. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
CarloLucibello
force-pushed
the
cl/dataloader-numworkers-guide-and-fix
branch
from
0dac018 to
36e7da5
Compare
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.
Enables and documents process-parallel data loading for
Dataset:MLUtils.DataLoader(ds; num_workers=N)spreadsgetobsoverNworker processes, each with its own CPython interpreter and GIL, so Python reads scale in a way the shared-GIL thread path (parallel=true) cannot. Works overmapobs/ObsView-wrapped datasets too.How it works
A
Datasetalready serializes across a process boundary viadatasets' own pickle — by reference to its on-disk Arrowcache_files, re-mmapped on the worker (data shared, not copied); in-memory datasets are materialized to a temporary Arrow dir once (fingerprint-cached). The requirement specific to distributed loading — that the pickle run on the GIL-holding main task, and that this package be loaded on the workers — is handled upstream in MLUtils (≥ 0.4.11), so it composes withmapobs/ObsViewat any nesting with nothing dataset-specific on the loader side.What's in the PR
num_workersworks out of the box.DataLoader(ds; num_workers=N)(and overmapobs/ObsView) loads in parallel worker processes. Requires MLUtils ≥ 0.4.11.__init__pinsJULIA_PYTHONCALL_EXEto the already-resolved interpreter when the env is CondaPkg-managed, so theNworkers don't each re-resolve the CondaPkg environment on startup (they'd otherwise serialize on CondaPkg's file lock and redo work the parent already did).num_workers), with two runnable doctests (batch structure and themapobs→(input, target)pattern) alongside Hub andnum_workersexamples.examples/flux_mnist.jl→ anexamples/flux_mnist/directory: a Flux + HuggingFaceDatasets MNIST script plus 1:1 and idiomatic PyTorch counterparts, and a README comparing data-loading strategies (serial / materialized / thread-parallel / process-parallel) across both stacks.@infowhen an in-memory dataset is first materialized to a temporary Arrow file for serialization.num_workersprocess-parallel loading.serialization.jlrefactor (_pickle_bytes/_unpickle_pyhelpers).[Unreleased]into[0.4.0]/[0.4.1]/[Unreleased], with thenum_workersentry under[Unreleased].🤖 Generated with Claude Code