From a026644c33abad9ed1c4303c77d10c6989a4eb45 Mon Sep 17 00:00:00 2001 From: Carlo Lucibello Date: Fri, 3 Jul 2026 09:07:40 +0200 Subject: [PATCH 1/8] Add process-parallel DataLoader (num_workers) support for Dataset MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `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) --- CHANGELOG.md | 13 +++++++- docs/src/guide.md | 60 +++++++++++++++++++++++++----------- src/HuggingFaceDatasets.jl | 5 +++ src/dataloader.jl | 60 ++++++++++++++++++++++++++++++++++++ src/serialization.jl | 24 ++++++++++----- test/Project.toml | 13 +++++--- test/serialization.jl | 63 ++++++++++++++++++++++++++++---------- 7 files changed, 191 insertions(+), 47 deletions(-) create mode 100644 src/dataloader.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index 65de54b..b7f1a07 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- `MLUtils.DataLoader(ds::Dataset; num_workers=N)` now works out of the box for process-parallel + loading. Passing a `Dataset` with `num_workers > 0` transparently wraps it in a new exported + `DistributedDataset`, which precomputes the `pickle` payload on the calling (GIL-holding) task + so the loader's background feeder ships bytes instead of calling Python off the GIL — which + segfaulted. An `@info` notes when an in-memory dataset is first materialized to a temporary + Arrow file for this. Requires MLUtils ≥ 0.4.10 (the first release with `num_workers`). + +## [0.4.0] - 2026-07-02 + This is a breaking release (0.3.x → 0.4.0). The headline change is that datasets are now returned in the `"julia"` format by default, so observations come back as native Julia values instead of raw Python objects. See **Breaking** below before upgrading. @@ -188,5 +198,6 @@ Baseline. Changes up to and including this release are recorded in the [git history](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/commits/main) and the [GitHub releases](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/releases). -[Unreleased]: https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/compare/v0.3.4...HEAD +[Unreleased]: https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/compare/v0.4.0...HEAD +[0.4.0]: https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/compare/v0.3.4...v0.4.0 [0.3.4]: https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/releases/tag/v0.3.4 diff --git a/docs/src/guide.md b/docs/src/guide.md index 527e158..1656b07 100644 --- a/docs/src/guide.md +++ b/docs/src/guide.md @@ -11,7 +11,8 @@ end This guide covers how the wrapper relates to the underlying Python `datasets` library, the `"julia"` format and the transform pipeline, array orientation and working with -images, and how to feed a dataset into a Julia data loader. +images, and how to feed a dataset into a Julia data loader — including process-parallel +(`num_workers`) loading that scales past CPython's GIL. The examples below build small datasets in memory with the [`Dataset`](@ref) constructor (which accepts a `Dict` or `NamedTuple` of columns) so that they are self-contained and @@ -422,7 +423,7 @@ numeric array**, not an `ImageCore` colorview. Combined with the axis reversal a while a ragged (variable-size) column falls back to a `Vector` of per-row arrays. That raw layout is what you want for feeding a model (see -[Integration with MLUtils and data loaders](@ref) and the MNIST example), but to *look* at +[Data loaders](@ref) and the MNIST example), but to *look* at an image you turn it back into a colorview and undo the transpose with `permutedims`. Two small helpers cover the common cases: @@ -467,11 +468,11 @@ The same `to_rgb` helper handles color datasets (e.g. `to_rgb(ds[1]["img"])` on and it round-trips exactly: the recovered array is pixel-for-pixel identical to the image `datasets` decodes on the Python side. -## Integration with MLUtils and data loaders +## Integration with MLUtils A [`Dataset`](@ref) implements the length/`getindex` interface expected by -[MLUtils.jl](https://github.com/JuliaML/MLUtils.jl), so `numobs`, `getobs`, `mapobs`, and -data loaders such as `Flux.DataLoader` work directly: +[MLUtils.jl](https://github.com/JuliaML/MLUtils.jl), so `numobs`, `getobs`, and `mapobs` +work directly: ```jldoctest guide julia> using MLUtils @@ -495,23 +496,46 @@ julia> getobs(mapped, 1:4) 40 ``` -Putting it together for MNIST: the default `"julia"` format already stacks the image -column into a `(W, H, N)` numeric array, so you `mapobs` a transform that rescales it to -`Float32` (plus `Flux.onehotbatch` for the labels), then hand the result to a -`Flux.DataLoader`: +Because of this, a `Dataset` can be handed straight to a data loader — see below. + +## Data Loaders + +Since a [`Dataset`](@ref) is a valid MLUtils data container, you can pass one straight to +`MLUtils.DataLoader` (re-exported by Flux as `Flux.DataLoader`). See the +[MLUtils docs](https://juliaml.github.io/MLUtils.jl/stable/) for the full set of options; the +points below are specific to a `Dataset`. ```julia-repl -julia> train_data = load_dataset("ylecun/mnist", split="train"); +julia> using MLUtils -julia> train_data = mapobs(mnist_transform, train_data)[:]; # lazily map, then materialize +julia> ds = load_dataset("ylecun/mnist", split="train"); # "julia" format by default -julia> train_loader = Flux.DataLoader(train_data; batchsize=128, shuffle=true); +julia> loader = DataLoader(ds; batchsize=128, shuffle=true); ``` -Materializing with `[:]` loads the whole (transformed) dataset into memory, which is -fastest for small datasets like MNIST. Dropping the `[:]` keeps loading on-the-fly, which -is slower per epoch but avoids holding everything in memory. - -This snippet needs the Hub and the Flux/ImageCore stack, so it is not run as a doctest; a -complete, runnable version — including `mnist_transform` and the training loop — lives in +For small datasets you can **materialize** into memory up front with `[:]` (faster per epoch, +no decoding during iteration); otherwise the loader reads **on the fly**, keeping memory flat. +A complete example — including the `mapobs` transform and training loop — lives in [`examples/flux_mnist.jl`](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/blob/main/examples/flux_mnist.jl). + +### Process-parallel loading with `num_workers` + +When loading on the fly, `getobs` calls into CPython, and CPython's global interpreter lock +(**GIL**) serializes those reads — so `parallel=true` (threads) gives ~1× on a `Dataset`. +Use `num_workers=N` instead (MLUtils v0.4.10 or newer): it spreads `getobs` over `N` worker +**processes** (via `Distributed`, mirroring PyTorch), each with its own interpreter and GIL, +so loading scales near-linearly. + +```julia-repl +julia> loader = DataLoader(ds; batchsize=128, shuffle=true, num_workers=4); + +julia> for batch in loader + x, y = batch["image"], batch["label"] # built by 4 workers in parallel + # ... training step ... + end +``` + +It works because a `Dataset` is `Serialization`-compatible +— it pickles by *reference* to its memory-mapped Arrow files, so workers re-mmap rather than +copy. This requires the default `"julia"` format (so `getobs` returns serializable arrays), a +serializable `jltransform` if you set one, and a shared filesystem across workers. diff --git a/src/HuggingFaceDatasets.jl b/src/HuggingFaceDatasets.jl index 6151234..c61cc1c 100644 --- a/src/HuggingFaceDatasets.jl +++ b/src/HuggingFaceDatasets.jl @@ -56,6 +56,11 @@ export concatenate_datasets, include("serialization.jl") +# Transparent `MLUtils.DataLoader(ds; num_workers=N)` support (precompute the worker payload). +include("dataloader.jl") +export DistributedDataset + +# `public` is a Julia 1.11+ keyword; `@compat` makes it a no-op on the supported 1.10. @compat public from_csv, from_json, from_parquet function __init__() diff --git a/src/dataloader.jl b/src/dataloader.jl new file mode 100644 index 0000000..cb9a88d --- /dev/null +++ b/src/dataloader.jl @@ -0,0 +1,60 @@ +# Transparent integration with `MLUtils.DataLoader` (re-exported by Flux) for process-parallel +# data loading over a `Dataset`. +# +# `DataLoader(ds; num_workers=N)` ships its data container to worker processes by serializing it +# from a background feeder task that does **not** hold the Python GIL; a `Dataset` serializes by +# calling `pickle`/PythonCall, which segfaults off the GIL-holding main task. The fix is to +# precompute the `pickle` bytes up front, on the main task, and ship *those* — pure Julia, safe +# from any thread. `DistributedDataset` is that prepared wrapper, and the `DataLoader(::Dataset)` +# method below installs it automatically when `num_workers > 0`. + +""" + DistributedDataset(ds::Dataset) + +A [`Dataset`](@ref) prepared to be sent to `Distributed` worker processes. Construction +precomputes — on the current, GIL-holding task — the `pickle` bytes that carry `ds` (its data, +by reference to the on-disk Arrow files, plus its Python format). Serializing the wrapper later +just ships those bytes and touches no Python, so it is safe from the GIL-less feeder task of a +process-parallel data loader. + +You rarely construct one directly: `MLUtils.DataLoader(ds; num_workers=N)` wraps `ds` in a +`DistributedDataset` for you. It forwards the `numobs`/`getobs` data-container interface to the +underlying dataset, so it is a drop-in stand-in for `ds`. +""" +struct DistributedDataset + ds::Dataset # valid on the process that built it; used for numobs/getobs there + bytes::Vector{UInt8} # precomputed pickle payload (data by reference + Python format) + jltransform # rides alongside the bytes +end + +DistributedDataset(ds::Dataset) = + DistributedDataset(ds, _pickle_bytes(ds), getfield(ds, :jltransform)) + +# Forward the data-container interface to the wrapped dataset (on the main process at +# construction, on the worker after `deserialize` has rebuilt a live `Dataset`). +MLUtils.numobs(dd::DistributedDataset) = numobs(dd.ds) +MLUtils.getobs(dd::DistributedDataset, i) = getobs(dd.ds, i) + +# Ship the precomputed bytes — pure Julia, safe from the GIL-less feeder thread — never a `Py`. +function Serialization.serialize(s::AbstractSerializer, dd::DistributedDataset) + Serialization.serialize_type(s, DistributedDataset) + Serialization.serialize(s, dd.bytes) + Serialization.serialize(s, dd.jltransform) + return nothing +end + +function Serialization.deserialize(s::AbstractSerializer, ::Type{DistributedDataset}) + bytes = Serialization.deserialize(s) + jltransform = Serialization.deserialize(s) + ds = set_jltransform!(Dataset(_unpickle_py(bytes)), jltransform) # worker holds the GIL here + return DistributedDataset(ds, bytes, jltransform) +end + +# Wrap in `DistributedDataset` for process-parallel loading so the feeder-thread `serialize` ships +# precomputed bytes instead of calling Python. Serial and thread-parallel (`parallel`) loaders +# serialize nothing on the main process, so they get the raw `ds`. `invoke` reaches the generic +# `DataLoader(::Any; …)` constructor (the wrapped data is not a `Dataset`, so no recursion). +function MLUtils.DataLoader(ds::Dataset; num_workers::Integer = 0, kws...) + data = num_workers > 0 ? DistributedDataset(ds) : ds + return invoke(MLUtils.DataLoader, Tuple{Any}, data; num_workers, kws...) +end diff --git a/src/serialization.jl b/src/serialization.jl index 4ff2dcd..8f0fbbc 100644 --- a/src/serialization.jl +++ b/src/serialization.jl @@ -35,6 +35,17 @@ using Serialization: AbstractSerializer const _SAVE_CACHE = Dict{String,String}() const _SAVE_LOCK = ReentrantLock() +# The `pickle` bytes that carry a dataset (by reference to its on-disk Arrow files) plus its +# Python format across a process boundary. Calling this touches Python, so it must run on a +# task holding the GIL (see `DistributedDataset`, which precomputes it on the main task). +_pickle_bytes(ds::Dataset) = pyconvert(Vector{UInt8}, pickle.dumps(_ondisk_py(ds))) + +# Reconstruct the `Py` (re-mmapping the referenced Arrow files) from bytes produced by +# `_pickle_bytes`. `pickle.loads` runs arbitrary code; safe here because the bytes come from our +# own `serialize` over a trusted `Distributed` cluster (no different from Distributed running +# serialized closures). Do not point it at untrusted bytes. +_unpickle_py(bytes) = pickle.loads(pybytes(bytes)) + # A `datasets.Dataset` `Py` backed by on-disk Arrow files, so `pickle` references the files # (zero-copy re-mmap on the worker) instead of embedding the in-memory buffer. On-disk # datasets (Hub / `load_from_disk` / `from_*`) already qualify; in-memory ones are @@ -44,7 +55,10 @@ function _ondisk_py(ds::Dataset) length(py.cache_files) == 0 || return py dir = lock(_SAVE_LOCK) do get!(_SAVE_CACHE, pyconvert(String, py._fingerprint)) do - d = mktempdir(); py.save_to_disk(d); d # auto-removed at process exit + d = mktempdir() # auto-removed at process exit + @info "HuggingFaceDatasets: writing in-memory dataset to a temporary Arrow file." path = d + py.save_to_disk(d) + d end end return getfield(load_from_disk(dir)::Dataset, :py) @@ -54,7 +68,7 @@ function Serialization.serialize(s::AbstractSerializer, ds::Dataset) Serialization.serialize_type(s, Dataset) # `pickle` carries the data (by reference to the on-disk Arrow files) AND the Python # format; the Julia transform is a separate field of the wrapper, so it rides alongside. - Serialization.serialize(s, pyconvert(Vector{UInt8}, pickle.dumps(_ondisk_py(ds)))) + Serialization.serialize(s, _pickle_bytes(ds)) Serialization.serialize(s, getfield(ds, :jltransform)) return nothing end @@ -62,9 +76,5 @@ end function Serialization.deserialize(s::AbstractSerializer, ::Type{Dataset}) bytes = Serialization.deserialize(s) jltransform = Serialization.deserialize(s) - # `pickle.loads` runs arbitrary code; safe here because `bytes` come from our own - # `serialize` over a trusted `Distributed` cluster (no different from Distributed running - # serialized closures). Do not point it at untrusted bytes. - py = pickle.loads(pybytes(bytes)) # re-mmaps the referenced Arrow files - return set_jltransform!(Dataset(py), jltransform) # Python format already restored by pickle + return set_jltransform!(Dataset(_unpickle_py(bytes)), jltransform) # format restored by pickle end diff --git a/test/Project.toml b/test/Project.toml index 76633f2..8e33182 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,17 +1,20 @@ -[sources] -HuggingFaceDatasets = {path = ".."} - [deps] +DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b" +HuggingFaceDatasets = "d94b9a45-fdf5-4270-b024-5cbb9ef7117d" ImageCore = "a09fc81d-aa75-5fe9-8630-4744c3626534" +MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d" -HuggingFaceDatasets = "d94b9a45-fdf5-4270-b024-5cbb9ef7117d" Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +[sources] +HuggingFaceDatasets = {path = ".."} + [compat] +DataStructures = "0.19" Distributed = "1" ImageCore = "0.9, 0.10" +MLUtils = "0.4.10" PythonCall = "0.9" Serialization = "1" - diff --git a/test/serialization.jl b/test/serialization.jl index 6f3e534..4a614e8 100644 --- a/test/serialization.jl +++ b/test/serialization.jl @@ -1,5 +1,7 @@ using Serialization: serialize, deserialize using Distributed +using MLUtils: DataLoader, getobs, numobs +import MLUtils # serialize → deserialize through an in-memory buffer (the same path `Distributed` uses to # ship a `Dataset` to a worker, minus the process hop). A worker read is exercised @@ -37,21 +39,50 @@ end @test length(HuggingFaceDatasets._SAVE_CACHE) == n0 + 1 end -# The actual cross-process path (serialize → ship to a worker → re-mmap → read there) needs -# an extra worker process, each spinning up its own Python — too heavy for CI. It guards the -# invariant that no `Py` crosses the boundary, which the in-process round-trips above cannot -# catch (a stray `Py` round-trips fine within a single interpreter but segfaults across one). -if !parse(Bool, get(ENV, "CI", "false")) - @testset "cross-process round-trip" begin - procs = addprocs(1; exeflags = "--project=$(dirname(Base.active_project()))") - try - @everywhere procs using HuggingFaceDatasets - ds = Dataset((; x = reshape(collect(1:8*20), 8, 20), label = collect(0:19))) - got = remotecall_fetch(d -> d[1:20], only(procs), ds) # ds serialized to the worker - @test got["x"] == ds[1:20]["x"] - @test got["label"] == ds[1:20]["label"] - finally - rmprocs(procs) - end +@testset "DistributedDataset round-trip" begin + # The feeder-safe wrapper installed by `DataLoader(ds; num_workers>0)`: it precomputes the + # pickle bytes on this task and serializes *those* (no `Py`, no Python call), so it survives + # being shipped from a GIL-less thread. Round-tripping must rebuild a working dataset. + ds = Dataset((; x = reshape(collect(1:8*20), 8, 20), label = collect(0:19))) + dd = DistributedDataset(ds) + @test dd isa DistributedDataset + @test numobs(dd) == 20 + r = roundtrip(dd) + @test r isa DistributedDataset + @test getobs(r, 1:20)["label"] == ds[1:20]["label"] + @test getobs(r, 1:20)["x"] == ds[1:20]["x"] +end + +# The real cross-process paths (serialize → ship to a worker → re-mmap → read there) each spin +# up a worker process with its own Python. They guard invariants the in-process round-trips +# above cannot: that no `Py` crosses the boundary (a stray `Py` round-trips fine within one +# interpreter but segfaults across one), and that materializing an in-memory dataset does not +# crash a concurrent loader. Worth the few extra seconds on CI. +@testset "cross-process round-trip" begin + procs = addprocs(1; exeflags = "--project=$(dirname(Base.active_project()))") + try + @everywhere procs using HuggingFaceDatasets + ds = Dataset((; x = reshape(collect(1:8*20), 8, 20), label = collect(0:19))) + got = remotecall_fetch(d -> d[1:20], only(procs), ds) # ds serialized to the worker + @test got["x"] == ds[1:20]["x"] + @test got["label"] == ds[1:20]["label"] + finally + rmprocs(procs) + end +end + +# End-to-end process-parallel loading of an *in-memory* dataset via MLUtils. This drives the +# full `DataLoader(...; num_workers=N)` path: the `DataLoader(::Dataset)` hook wraps `ds` in a +# `DistributedDataset` (precomputing the pickle bytes on the main task) so the background feeder +# ships those bytes instead of calling Python off the GIL-holding thread, which used to segfault. +@testset "num_workers DataLoader over in-memory dataset" begin + ds = Dataset((; x = reshape(collect(1:8*40), 8, 40), label = collect(0:39))) + loader = DataLoader(ds; batchsize = 10, num_workers = 2) + try + batches = collect(loader) # iterates ⇒ ships ds to 2 workers + @test length(batches) == 4 + @test sort(reduce(vcat, [b["label"] for b in batches])) == collect(0:39) + finally + MLUtils.close_dataloader_pool() end end From 265acdf7436ae2f61228c5915cbb43d4a1fd0dca Mon Sep 17 00:00:00 2001 From: Carlo Lucibello Date: Fri, 3 Jul 2026 09:07:46 +0200 Subject: [PATCH 2/8] docs: dedicated Data loaders section in the guide 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) --- docs/src/api.md | 6 ++++++ docs/src/guide.md | 8 +++++--- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/docs/src/api.md b/docs/src/api.md index 63fa0f3..76a9b34 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -47,6 +47,12 @@ concatenate_datasets interleave_datasets ``` +## Data loading + +```@docs +DistributedDataset +``` + ## Formats and transforms ```@docs diff --git a/docs/src/guide.md b/docs/src/guide.md index 1656b07..d61e23a 100644 --- a/docs/src/guide.md +++ b/docs/src/guide.md @@ -535,7 +535,9 @@ julia> for batch in loader end ``` -It works because a `Dataset` is `Serialization`-compatible -— it pickles by *reference* to its memory-mapped Arrow files, so workers re-mmap rather than -copy. This requires the default `"julia"` format (so `getobs` returns serializable arrays), a +Passing a `Dataset` to `DataLoader` with `num_workers > 0` transparently wraps it in a +[`DistributedDataset`](@ref), which precomputes — on the calling task — the payload each worker +needs (an `@info` notes when an in-memory dataset is first materialized to a temporary Arrow file +for this). Workers re-mmap the dataset's Arrow files by *reference* rather than copying it, so +this requires the default `"julia"` format (so `getobs` returns serializable arrays), a serializable `jltransform` if you set one, and a shared filesystem across workers. From 203232808f1034e58b2c7fed4d142cf77a6a3fd1 Mon Sep 17 00:00:00 2001 From: Carlo Lucibello Date: Fri, 3 Jul 2026 12:40:01 +0200 Subject: [PATCH 3/8] Rely on MLUtils for num_workers serialization; drop DistributedDataset 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) --- CHANGELOG.md | 12 ++++---- docs/src/api.md | 6 ---- docs/src/guide.md | 15 +++++----- src/HuggingFaceDatasets.jl | 4 --- src/dataloader.jl | 60 -------------------------------------- src/serialization.jl | 3 +- test/Project.toml | 2 +- test/serialization.jl | 24 ++++----------- 8 files changed, 22 insertions(+), 104 deletions(-) delete mode 100644 src/dataloader.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index b7f1a07..4c9da59 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,12 +8,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added -- `MLUtils.DataLoader(ds::Dataset; num_workers=N)` now works out of the box for process-parallel - loading. Passing a `Dataset` with `num_workers > 0` transparently wraps it in a new exported - `DistributedDataset`, which precomputes the `pickle` payload on the calling (GIL-holding) task - so the loader's background feeder ships bytes instead of calling Python off the GIL — which - segfaulted. An `@info` notes when an in-memory dataset is first materialized to a temporary - Arrow file for this. Requires MLUtils ≥ 0.4.10 (the first release with `num_workers`). +- `MLUtils.DataLoader(ds::Dataset; num_workers=N)` works out of the box for process-parallel + loading, including over `mapobs`/`ObsView`-wrapped datasets. The GIL-safe serialization it + relies on — running the dataset's `pickle` on the main task and loading this package on the + workers — is handled upstream by MLUtils, so no wrapper type is exported here. An `@info` notes + when an in-memory dataset is first materialized to a temporary Arrow file for this. Requires + MLUtils ≥ 0.4.11. ## [0.4.0] - 2026-07-02 diff --git a/docs/src/api.md b/docs/src/api.md index 76a9b34..63fa0f3 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -47,12 +47,6 @@ concatenate_datasets interleave_datasets ``` -## Data loading - -```@docs -DistributedDataset -``` - ## Formats and transforms ```@docs diff --git a/docs/src/guide.md b/docs/src/guide.md index d61e23a..591a68a 100644 --- a/docs/src/guide.md +++ b/docs/src/guide.md @@ -522,7 +522,7 @@ A complete example — including the `mapobs` transform and training loop — li When loading on the fly, `getobs` calls into CPython, and CPython's global interpreter lock (**GIL**) serializes those reads — so `parallel=true` (threads) gives ~1× on a `Dataset`. -Use `num_workers=N` instead (MLUtils v0.4.10 or newer): it spreads `getobs` over `N` worker +Use `num_workers=N` instead (MLUtils v0.4.11 or newer): it spreads `getobs` over `N` worker **processes** (via `Distributed`, mirroring PyTorch), each with its own interpreter and GIL, so loading scales near-linearly. @@ -535,9 +535,10 @@ julia> for batch in loader end ``` -Passing a `Dataset` to `DataLoader` with `num_workers > 0` transparently wraps it in a -[`DistributedDataset`](@ref), which precomputes — on the calling task — the payload each worker -needs (an `@info` notes when an in-memory dataset is first materialized to a temporary Arrow file -for this). Workers re-mmap the dataset's Arrow files by *reference* rather than copying it, so -this requires the default `"julia"` format (so `getobs` returns serializable arrays), a -serializable `jltransform` if you set one, and a shared filesystem across workers. +Under `num_workers > 0` the loader serializes `ds` on the calling task — so the Python work +happens where the GIL makes it safe — and ships each worker the payload it needs; an `@info` +notes when an in-memory dataset is first materialized to a temporary Arrow file for this. Workers +re-mmap the dataset's Arrow files by *reference* rather than copying it, so this requires the +default `"julia"` format (so `getobs` returns serializable arrays), a serializable `jltransform` +if you set one, and a shared filesystem across workers. Wrapping the dataset in `mapobs`/`ObsView` +before the loader is supported too. diff --git a/src/HuggingFaceDatasets.jl b/src/HuggingFaceDatasets.jl index c61cc1c..2a26859 100644 --- a/src/HuggingFaceDatasets.jl +++ b/src/HuggingFaceDatasets.jl @@ -56,10 +56,6 @@ export concatenate_datasets, include("serialization.jl") -# Transparent `MLUtils.DataLoader(ds; num_workers=N)` support (precompute the worker payload). -include("dataloader.jl") -export DistributedDataset - # `public` is a Julia 1.11+ keyword; `@compat` makes it a no-op on the supported 1.10. @compat public from_csv, from_json, from_parquet diff --git a/src/dataloader.jl b/src/dataloader.jl deleted file mode 100644 index cb9a88d..0000000 --- a/src/dataloader.jl +++ /dev/null @@ -1,60 +0,0 @@ -# Transparent integration with `MLUtils.DataLoader` (re-exported by Flux) for process-parallel -# data loading over a `Dataset`. -# -# `DataLoader(ds; num_workers=N)` ships its data container to worker processes by serializing it -# from a background feeder task that does **not** hold the Python GIL; a `Dataset` serializes by -# calling `pickle`/PythonCall, which segfaults off the GIL-holding main task. The fix is to -# precompute the `pickle` bytes up front, on the main task, and ship *those* — pure Julia, safe -# from any thread. `DistributedDataset` is that prepared wrapper, and the `DataLoader(::Dataset)` -# method below installs it automatically when `num_workers > 0`. - -""" - DistributedDataset(ds::Dataset) - -A [`Dataset`](@ref) prepared to be sent to `Distributed` worker processes. Construction -precomputes — on the current, GIL-holding task — the `pickle` bytes that carry `ds` (its data, -by reference to the on-disk Arrow files, plus its Python format). Serializing the wrapper later -just ships those bytes and touches no Python, so it is safe from the GIL-less feeder task of a -process-parallel data loader. - -You rarely construct one directly: `MLUtils.DataLoader(ds; num_workers=N)` wraps `ds` in a -`DistributedDataset` for you. It forwards the `numobs`/`getobs` data-container interface to the -underlying dataset, so it is a drop-in stand-in for `ds`. -""" -struct DistributedDataset - ds::Dataset # valid on the process that built it; used for numobs/getobs there - bytes::Vector{UInt8} # precomputed pickle payload (data by reference + Python format) - jltransform # rides alongside the bytes -end - -DistributedDataset(ds::Dataset) = - DistributedDataset(ds, _pickle_bytes(ds), getfield(ds, :jltransform)) - -# Forward the data-container interface to the wrapped dataset (on the main process at -# construction, on the worker after `deserialize` has rebuilt a live `Dataset`). -MLUtils.numobs(dd::DistributedDataset) = numobs(dd.ds) -MLUtils.getobs(dd::DistributedDataset, i) = getobs(dd.ds, i) - -# Ship the precomputed bytes — pure Julia, safe from the GIL-less feeder thread — never a `Py`. -function Serialization.serialize(s::AbstractSerializer, dd::DistributedDataset) - Serialization.serialize_type(s, DistributedDataset) - Serialization.serialize(s, dd.bytes) - Serialization.serialize(s, dd.jltransform) - return nothing -end - -function Serialization.deserialize(s::AbstractSerializer, ::Type{DistributedDataset}) - bytes = Serialization.deserialize(s) - jltransform = Serialization.deserialize(s) - ds = set_jltransform!(Dataset(_unpickle_py(bytes)), jltransform) # worker holds the GIL here - return DistributedDataset(ds, bytes, jltransform) -end - -# Wrap in `DistributedDataset` for process-parallel loading so the feeder-thread `serialize` ships -# precomputed bytes instead of calling Python. Serial and thread-parallel (`parallel`) loaders -# serialize nothing on the main process, so they get the raw `ds`. `invoke` reaches the generic -# `DataLoader(::Any; …)` constructor (the wrapped data is not a `Dataset`, so no recursion). -function MLUtils.DataLoader(ds::Dataset; num_workers::Integer = 0, kws...) - data = num_workers > 0 ? DistributedDataset(ds) : ds - return invoke(MLUtils.DataLoader, Tuple{Any}, data; num_workers, kws...) -end diff --git a/src/serialization.jl b/src/serialization.jl index 8f0fbbc..35603b0 100644 --- a/src/serialization.jl +++ b/src/serialization.jl @@ -37,7 +37,8 @@ const _SAVE_LOCK = ReentrantLock() # The `pickle` bytes that carry a dataset (by reference to its on-disk Arrow files) plus its # Python format across a process boundary. Calling this touches Python, so it must run on a -# task holding the GIL (see `DistributedDataset`, which precomputes it on the main task). +# task holding the GIL; the process-parallel `DataLoader` path serializes a `Dataset` on the +# main (GIL-holding) task for exactly this reason, so nothing calls Python off the GIL. _pickle_bytes(ds::Dataset) = pyconvert(Vector{UInt8}, pickle.dumps(_ondisk_py(ds))) # Reconstruct the `Py` (re-mmapping the referenced Arrow files) from bytes produced by diff --git a/test/Project.toml b/test/Project.toml index 8e33182..e6dea99 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -15,6 +15,6 @@ HuggingFaceDatasets = {path = ".."} DataStructures = "0.19" Distributed = "1" ImageCore = "0.9, 0.10" -MLUtils = "0.4.10" +MLUtils = "0.4.11" PythonCall = "0.9" Serialization = "1" diff --git a/test/serialization.jl b/test/serialization.jl index 4a614e8..2e2a4ef 100644 --- a/test/serialization.jl +++ b/test/serialization.jl @@ -1,6 +1,6 @@ using Serialization: serialize, deserialize using Distributed -using MLUtils: DataLoader, getobs, numobs +using MLUtils: DataLoader import MLUtils # serialize → deserialize through an in-memory buffer (the same path `Distributed` uses to @@ -39,20 +39,6 @@ end @test length(HuggingFaceDatasets._SAVE_CACHE) == n0 + 1 end -@testset "DistributedDataset round-trip" begin - # The feeder-safe wrapper installed by `DataLoader(ds; num_workers>0)`: it precomputes the - # pickle bytes on this task and serializes *those* (no `Py`, no Python call), so it survives - # being shipped from a GIL-less thread. Round-tripping must rebuild a working dataset. - ds = Dataset((; x = reshape(collect(1:8*20), 8, 20), label = collect(0:19))) - dd = DistributedDataset(ds) - @test dd isa DistributedDataset - @test numobs(dd) == 20 - r = roundtrip(dd) - @test r isa DistributedDataset - @test getobs(r, 1:20)["label"] == ds[1:20]["label"] - @test getobs(r, 1:20)["x"] == ds[1:20]["x"] -end - # The real cross-process paths (serialize → ship to a worker → re-mmap → read there) each spin # up a worker process with its own Python. They guard invariants the in-process round-trips # above cannot: that no `Py` crosses the boundary (a stray `Py` round-trips fine within one @@ -71,10 +57,10 @@ end end end -# End-to-end process-parallel loading of an *in-memory* dataset via MLUtils. This drives the -# full `DataLoader(...; num_workers=N)` path: the `DataLoader(::Dataset)` hook wraps `ds` in a -# `DistributedDataset` (precomputing the pickle bytes on the main task) so the background feeder -# ships those bytes instead of calling Python off the GIL-holding thread, which used to segfault. +# End-to-end process-parallel loading of an *in-memory* dataset via MLUtils. This drives the full +# `DataLoader(...; num_workers=N)` path: MLUtils serializes `ds` on the main task (so the pickle +# runs under the GIL) and loads this package on the workers, so the background feeder never calls +# Python off the GIL-holding thread — which used to segfault. Requires MLUtils ≥ 0.4.11. @testset "num_workers DataLoader over in-memory dataset" begin ds = Dataset((; x = reshape(collect(1:8*40), 8, 40), label = collect(0:39))) loader = DataLoader(ds; batchsize = 10, num_workers = 2) From ede10d241486d794db623a30d59fd2b1d435997c Mon Sep 17 00:00:00 2001 From: Carlo Lucibello Date: Sat, 4 Jul 2026 07:46:04 +0200 Subject: [PATCH 4/8] Fix docs @ref, trim Data Loaders guide, num_workers MNIST example - 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) --- docs/src/guide.md | 44 +++++---------- examples/flux_mnist.jl | 71 ------------------------ examples/{ => flux_mnist}/Project.toml | 9 +++ examples/flux_mnist/main.jl | 77 ++++++++++++++++++++++++++ 4 files changed, 100 insertions(+), 101 deletions(-) delete mode 100644 examples/flux_mnist.jl rename examples/{ => flux_mnist}/Project.toml (54%) create mode 100644 examples/flux_mnist/main.jl diff --git a/docs/src/guide.md b/docs/src/guide.md index 591a68a..7e18729 100644 --- a/docs/src/guide.md +++ b/docs/src/guide.md @@ -423,7 +423,7 @@ numeric array**, not an `ImageCore` colorview. Combined with the axis reversal a while a ragged (variable-size) column falls back to a `Vector` of per-row arrays. That raw layout is what you want for feeding a model (see -[Data loaders](@ref) and the MNIST example), but to *look* at +[Data Loaders](@ref) and the MNIST example), but to *look* at an image you turn it back into a colorview and undo the transpose with `permutedims`. Two small helpers cover the common cases: @@ -500,10 +500,11 @@ Because of this, a `Dataset` can be handed straight to a data loader — see bel ## Data Loaders -Since a [`Dataset`](@ref) is a valid MLUtils data container, you can pass one straight to -`MLUtils.DataLoader` (re-exported by Flux as `Flux.DataLoader`). See the -[MLUtils docs](https://juliaml.github.io/MLUtils.jl/stable/) for the full set of options; the -points below are specific to a `Dataset`. +A [`Dataset`](@ref) is a valid MLUtils data container, so you can hand one straight to +`MLUtils.DataLoader` (re-exported by Flux as `Flux.DataLoader`), including through a +`mapobs`/`ObsView` transform. The loader's own options are covered in the +[MLUtils Data Loaders guide](https://juliaml.github.io/MLUtils.jl/stable/guide/iteration/); +the notes below are specific to a `Dataset`. ```julia-repl julia> using MLUtils @@ -515,30 +516,13 @@ julia> loader = DataLoader(ds; batchsize=128, shuffle=true); For small datasets you can **materialize** into memory up front with `[:]` (faster per epoch, no decoding during iteration); otherwise the loader reads **on the fly**, keeping memory flat. -A complete example — including the `mapobs` transform and training loop — lives in -[`examples/flux_mnist.jl`](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/blob/main/examples/flux_mnist.jl). -### Process-parallel loading with `num_workers` +Reading on the fly calls into CPython, and CPython's global interpreter lock (**GIL**) +serializes those reads, so thread-based `parallel=true` gives ~1× on a `Dataset`. Use +`num_workers=N` instead (needs MLUtils ≥ 0.4.11): it spreads `getobs` over `N` worker +**processes**, each with its own interpreter and GIL, so loading scales past the GIL. This +works because a `Dataset` serializes by *reference* to its on-disk Arrow files (an in-memory +dataset is first materialized to a temporary Arrow file), so it needs a serializable `jltransform` if you set one. See the MLUtils guide for the mechanics and tradeoffs. -When loading on the fly, `getobs` calls into CPython, and CPython's global interpreter lock -(**GIL**) serializes those reads — so `parallel=true` (threads) gives ~1× on a `Dataset`. -Use `num_workers=N` instead (MLUtils v0.4.11 or newer): it spreads `getobs` over `N` worker -**processes** (via `Distributed`, mirroring PyTorch), each with its own interpreter and GIL, -so loading scales near-linearly. - -```julia-repl -julia> loader = DataLoader(ds; batchsize=128, shuffle=true, num_workers=4); - -julia> for batch in loader - x, y = batch["image"], batch["label"] # built by 4 workers in parallel - # ... training step ... - end -``` - -Under `num_workers > 0` the loader serializes `ds` on the calling task — so the Python work -happens where the GIL makes it safe — and ships each worker the payload it needs; an `@info` -notes when an in-memory dataset is first materialized to a temporary Arrow file for this. Workers -re-mmap the dataset's Arrow files by *reference* rather than copying it, so this requires the -default `"julia"` format (so `getobs` returns serializable arrays), a serializable `jltransform` -if you set one, and a shared filesystem across workers. Wrapping the dataset in `mapobs`/`ObsView` -before the loader is supported too. +A complete example — `mapobs` transform, `num_workers`, and a training loop — lives in +[`examples/flux_mnist/main.jl`](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/blob/main/examples/flux_mnist/main.jl). diff --git a/examples/flux_mnist.jl b/examples/flux_mnist.jl deleted file mode 100644 index b16b8a0..0000000 --- a/examples/flux_mnist.jl +++ /dev/null @@ -1,71 +0,0 @@ -using Flux, Zygote -using Random, Statistics -using Flux.Losses: logitcrossentropy -using Flux: onecold -using HuggingFaceDatasets -using MLUtils -# using ProfileView, BenchmarkTools - -function mnist_transform(batch) - # under the "julia" (= numpy) format the image column is already a stacked - # (W, H, N) UInt8 array, so just rescale to Float32 in [0, 1] - image = Float32.(batch["image"]) ./ 255f0 - label = Flux.onehotbatch(batch["label"], 0:9) - return (; image, label) -end - -function loss_and_accuracy(data_loader, model, device) - acc = 0 - ls = 0.0f0 - num = 0 - for (x, y) in data_loader - x, y = x |> device, y |> device - ŷ = model(x) - ls += logitcrossentropy(ŷ, y, agg=sum) - acc += sum(onecold(ŷ) .== onecold(y)) - num += size(x)[end] - end - return ls / num, acc / num -end - -function train(epochs) - batchsize = 128 - nhidden = 100 - device = cpu - - train_data = load_dataset("ylecun/mnist", split="train").with_format("julia") - test_data = load_dataset("ylecun/mnist", split="test").with_format("julia") - train_data = mapobs(mnist_transform, train_data)[:] # lazy apply transform then materialize - test_data = mapobs(mnist_transform, test_data)[:] - - train_loader = Flux.DataLoader(train_data; batchsize, shuffle=true) - test_loader = Flux.DataLoader(test_data; batchsize) - - model = Chain([Flux.flatten, - Dense(28*28, nhidden, relu), - Dense(nhidden, nhidden, relu), - Dense(nhidden, 10)]) |> device - - opt = Flux.setup(AdamW(1e-3), model) - - function report(epoch) - train_loss, train_acc = loss_and_accuracy(train_loader, model, device) - test_loss, test_acc = loss_and_accuracy(test_loader, model, device) - r(x) = round(x, digits=3) - r(x::Int) = x - @info map(r, (; epoch, train_loss, train_acc, test_loss, test_acc)) - end - - report(0) - @time for epoch in 1:epochs - for (x, y) in train_loader - x, y = x |> device, y |> device - loss, grads = withgradient(model -> logitcrossentropy(model(x), y), model) - Flux.update!(opt, model, grads[1]) - end - report(epoch) - end -end - -@time train(2) # 8s on a m1 pro with in-memory loading - # 20s on-the-fly loading diff --git a/examples/Project.toml b/examples/flux_mnist/Project.toml similarity index 54% rename from examples/Project.toml rename to examples/flux_mnist/Project.toml index 78c8105..f59c105 100644 --- a/examples/Project.toml +++ b/examples/flux_mnist/Project.toml @@ -1,6 +1,15 @@ [deps] +DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" HuggingFaceDatasets = "d94b9a45-fdf5-4270-b024-5cbb9ef7117d" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[sources] +HuggingFaceDatasets = {path = "../.."} + +[compat] +MLUtils = "0.4" diff --git a/examples/flux_mnist/main.jl b/examples/flux_mnist/main.jl new file mode 100644 index 0000000..338f7d6 --- /dev/null +++ b/examples/flux_mnist/main.jl @@ -0,0 +1,77 @@ +using Flux, Zygote +using Random, Statistics +using Flux.Losses: logitcrossentropy +using Flux: onecold, onehotbatch +using HuggingFaceDatasets +using MLUtils +# using ProfileView, BenchmarkTools + +function mnist_transform(batch) + # under the "julia" (= numpy) format the image column is already a stacked + # (W, H, N) UInt8 array, so just rescale to Float32 in [0, 1] + image = Float32.(batch["image"]) ./ 255f0 + return (; image, label = batch["label"]) +end + +function loss_and_accuracy(data_loader, model, device) + acc = 0 + ls = 0.0f0 + num = 0 + for (x, y) in data_loader + x = x |> device + yoh = onehotbatch(y, 0:9) |> device + ŷ = model(x) + ls += logitcrossentropy(ŷ, yoh, agg=sum) + acc += sum(onecold(ŷ, 0:9) .== y) + num += length(y) + end + return ls / num, acc / num +end + +# `num_workers = 0` loads on the main process; `num_workers > 0` spreads each batch's +# `getobs` (and the CPython read it triggers) over that many worker processes, sidestepping +# the GIL. MLUtils spawns the workers on demand under the current `--project`. +function train(; epochs=2, num_workers=0) + batchsize = 128 + nhidden = 100 + device = cpu + + train_data = load_dataset("ylecun/mnist", split="train") + test_data = load_dataset("ylecun/mnist", split="test") + # apply the transform lazily so it runs per batch during iteration (on the workers when + # `num_workers > 0`); `mapobs`/`ObsView`-wrapped datasets compose with `num_workers` + train_data = mapobs(mnist_transform, train_data) + test_data = mapobs(mnist_transform, test_data) + + train_loader = Flux.DataLoader(train_data; batchsize, shuffle=true, num_workers) + test_loader = Flux.DataLoader(test_data; batchsize, num_workers) + + model = Chain([Flux.flatten, + Dense(28*28, nhidden, relu), + Dense(nhidden, nhidden, relu), + Dense(nhidden, 10)]) |> device + + opt = Flux.setup(AdamW(1e-3), model) + + function report(epoch) + train_loss, train_acc = loss_and_accuracy(train_loader, model, device) + test_loss, test_acc = loss_and_accuracy(test_loader, model, device) + r(x) = round(x, digits=3) + r(x::Int) = x + @info map(r, (; epoch, train_loss, train_acc, test_loss, test_acc)) + end + + report(0) + @time for epoch in 1:epochs + for (x, y) in train_loader + x = x |> device + yoh = onehotbatch(y, 0:9) |> device + loss, grads = withgradient(m -> logitcrossentropy(m(x), yoh), model) + Flux.update!(opt, model, grads[1]) + end + report(epoch) + end +end + +# @time train(; epochs=2, num_workers=0) +@time train(; epochs=2, num_workers=4) From a7d60288a69402e52f18e57540d6b2c45baa80a0 Mon Sep 17 00:00:00 2001 From: Carlo Lucibello Date: Mon, 6 Jul 2026 09:34:29 +0200 Subject: [PATCH 5/8] Pin JULIA_PYTHONCALL_EXE so num_workers workers skip CondaPkg re-resolve 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 --- src/HuggingFaceDatasets.jl | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/HuggingFaceDatasets.jl b/src/HuggingFaceDatasets.jl index 2a26859..9fbfc42 100644 --- a/src/HuggingFaceDatasets.jl +++ b/src/HuggingFaceDatasets.jl @@ -70,6 +70,17 @@ function __init__() PythonCall.pycopy!(np, pyimport("numpy")) PythonCall.pycopy!(pycopy, pyimport("copy")) PythonCall.pycopy!(pickle, pyimport("pickle")) + + # Pin the already-resolved Python interpreter for any child processes. + # + # The recommended way to parallelize loading past the CPython GIL is + # `DataLoader(dataset; num_workers=N)`, which spawns `N` `Distributed` worker + # processes that each `using PythonCall`. Without the following, every worker independently + # *re-resolves* the CondaPkg environment on startup, in lockstep: they serialize on + # CondaPkg's file lock and redo work this parent process already did. + if PythonCall.C.CTX.which === :CondaPkg + get!(ENV, "JULIA_PYTHONCALL_EXE", PythonCall.python_executable_path()::String) + end end end # module From c0173a0af4c434185a708f414c0ba15997b10561 Mon Sep 17 00:00:00 2001 From: Carlo Lucibello Date: Mon, 6 Jul 2026 09:34:40 +0200 Subject: [PATCH 6/8] examples: flux_mnist data-loading benchmark, rename, doc links - 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 --- .gitignore | 3 +- docs/src/guide.md | 4 +- examples/flux_mnist/Project.toml | 3 +- examples/flux_mnist/README.md | 80 ++++++++++ .../flux_mnist/{main.jl => flux_mnist.jl} | 44 +++-- examples/flux_mnist/pytorch_mnist.py | 135 ++++++++++++++++ .../flux_mnist/pytorch_mnist_idiomatic.py | 151 ++++++++++++++++++ 7 files changed, 402 insertions(+), 18 deletions(-) create mode 100644 examples/flux_mnist/README.md rename examples/flux_mnist/{main.jl => flux_mnist.jl} (62%) create mode 100644 examples/flux_mnist/pytorch_mnist.py create mode 100644 examples/flux_mnist/pytorch_mnist_idiomatic.py diff --git a/.gitignore b/.gitignore index 32ccdc0..12ec012 100644 --- a/.gitignore +++ b/.gitignore @@ -7,4 +7,5 @@ Manifest.toml .vscode .ipynb_checkpoints LocalPreferences.toml -.claude/ \ No newline at end of file +.claude/ +__pycache__/ diff --git a/docs/src/guide.md b/docs/src/guide.md index 7e18729..c38cf35 100644 --- a/docs/src/guide.md +++ b/docs/src/guide.md @@ -293,7 +293,7 @@ Note that the format and the custom transform share the same slot: setting the ` format installs `py2jl` as the transform, and `with_jltransform` then replaces it — which is why the transform above calls `py2jl` itself. For layering additional per-batch processing on top of the `"julia"` format, prefer `MLUtils.mapobs` (see below), as in the -[`examples/flux_mnist.jl`](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/blob/main/examples/flux_mnist.jl) +[`examples/flux_mnist/flux_mnist.jl`](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/blob/main/examples/flux_mnist/flux_mnist.jl) script. The order of operations when you index is: @@ -525,4 +525,4 @@ works because a `Dataset` serializes by *reference* to its on-disk Arrow files ( dataset is first materialized to a temporary Arrow file), so it needs a serializable `jltransform` if you set one. See the MLUtils guide for the mechanics and tradeoffs. A complete example — `mapobs` transform, `num_workers`, and a training loop — lives in -[`examples/flux_mnist/main.jl`](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/blob/main/examples/flux_mnist/main.jl). +[`examples/flux_mnist/flux_mnist.jl`](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/blob/main/examples/flux_mnist/flux_mnist.jl). diff --git a/examples/flux_mnist/Project.toml b/examples/flux_mnist/Project.toml index f59c105..eeed48e 100644 --- a/examples/flux_mnist/Project.toml +++ b/examples/flux_mnist/Project.toml @@ -12,4 +12,5 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" HuggingFaceDatasets = {path = "../.."} [compat] -MLUtils = "0.4" +DataStructures = "0.19" +MLUtils = "0.4" \ No newline at end of file diff --git a/examples/flux_mnist/README.md b/examples/flux_mnist/README.md new file mode 100644 index 0000000..63af6db --- /dev/null +++ b/examples/flux_mnist/README.md @@ -0,0 +1,80 @@ +# Flux MNIST — data-loading benchmark + +Trains a small MLP on MNIST pulled from the HuggingFace `datasets` library, comparing +**data-loading strategies** (on-the-fly vs. materialized, thread- vs. process-parallel) in +Flux+HuggingFaceDatasets.jl against equivalent and idiomatic PyTorch versions. + +## Files + +| file | stack | notes | +| --- | --- | --- | +| [`flux_mnist.jl`](flux_mnist.jl) | Flux + HuggingFaceDatasets.jl | `julia --project=. -t4 flux_mnist.jl` | +| [`pytorch_mnist.py`](pytorch_mnist.py) | PyTorch + `datasets` | 1:1 port of `flux_mnist.jl`; `uv run pytorch_mnist.py` | +| [`pytorch_mnist_idiomatic.py`](pytorch_mnist_idiomatic.py) | PyTorch + `datasets` | idiomatic HF patterns (`with_transform`, `.map`); `uv run pytorch_mnist_idiomatic.py` | + +## Setup + +Same model and hyperparameters everywhere: MLP `784 → 100 → 100 → 10`, `AdamW(1e-3)`, +cross-entropy, batch size 128, **4 epochs**, **CPU only**, MNIST (`ylecun/mnist`) via HF +`datasets`. Timings are single-run wall-clock on an **Apple M1 Pro** — indicative, not +rigorous (expect ±10–20% run to run). A warm-up epoch precedes the Julia timings to exclude +Julia's one-time JIT compilation; PyTorch runs eagerly (no comparable compile step) so needs +none. Both then measure steady-state per-epoch compute — imports and process startup sit +outside the timed region either way — so the tables are directly comparable. + +## Results + +### Julia — Flux + HuggingFaceDatasets.jl (`-t4`) + +| config | data loading | time | +| --- | --- | ---: | +| Serial | on-the-fly, `num_workers=0` | 13.5 s | +| Serial Materialized | in-memory `[:]`, `num_workers=0` | 6.4 s | +| Parallel Materialized | in-memory, `parallel=true` (threads) | 9.1 s | +| Distributed | on-the-fly, `num_workers=4` (processes) | 33.0 s | + +A warm-up epoch runs first, so these exclude Julia's JIT compilation. The serial path is +then fully warmed; the `parallel`/`num_workers` paths still compile their first call, and +Distributed also pays worker-process startup — real costs for a short job, left in the numbers. + +### PyTorch — 1:1 port (`pytorch_mnist.py`) + +| config | data loading | time | +| --- | --- | ---: | +| Serial | on-the-fly, `num_workers=0` | 62.8 s | +| Serial Materialized | in-memory tensors, `num_workers=0` | 16.8 s | +| Parallel Materialized | in-memory, `num_workers=4` (processes) | 22.2 s | +| Distributed | on-the-fly, `num_workers=4` (processes) | 66.6 s | + +### PyTorch — idiomatic HF (`pytorch_mnist_idiomatic.py`) + +| config | data loading | time | +| --- | --- | ---: | +| Lazy `with_transform` | on-the-fly, `num_workers=0` | 34.1 s | +| Lazy `with_transform` | on-the-fly, `num_workers=4` | 55.1 s | +| Cached `.map` (torch format) | preprocessed to Arrow, `num_workers=0` | 291.4 s† | +| Cached `.map` (torch format) | preprocessed to Arrow, `num_workers=4` | 128.3 s | + +†includes the one-time `.map` decode+cache; the `num_workers=4` row reuses that cache. + +## Takeaways + +- **Materializing into memory is the biggest win** for a dataset this small — ~2× in Julia + (13.5 → 6.4 s) and ~4× in PyTorch (62.8 → 16.8 s), because it drops the per-batch CPython + decode entirely. +- **Parallel loading does *not* pay off here.** The MLP is tiny, so multiprocess + pickling/IPC overhead outweighs the parallel-decode benefit — `num_workers=4` was *slower* + than `num_workers=0` in every PyTorch case, and Julia's `Distributed` (33.0 s) lost to its + serial on-the-fly (13.5 s). `num_workers`/process parallelism is for when `getobs` is the + bottleneck (large images, heavy decode), not toy workloads. + - Julia's **thread**-based `parallel=true` over *materialized* data (9.1 s) is cheap + (shared memory, needs `-t>1`) but pointless here — there's nothing to parallelize once the + data is plain in-memory arrays. PyTorch has no thread analog (the GIL); its only knob is + `num_workers` (processes). +- **Idiomatic PyTorch is competitive but has a trap.** Lazy `with_transform` (34.1 s) even + beats the 1:1 Serial port (62.8 s), but naively `.map`-ing decoded float tensors into Arrow + and reading them back via the torch formatter is *much* slower (128–291 s) — for images, + prefer lazy transforms or materialize to plain tensors, don't cache decoded floats. +- **Julia came out ahead of PyTorch** on the like-for-like paths (on-the-fly 13.5 vs 62.8 s; + materialized 6.4 vs 16.8 s) — a fair comparison, since both are steady-state per-epoch + compute (Julia after a warm-up, PyTorch eager with no compile step to warm away). diff --git a/examples/flux_mnist/main.jl b/examples/flux_mnist/flux_mnist.jl similarity index 62% rename from examples/flux_mnist/main.jl rename to examples/flux_mnist/flux_mnist.jl index 338f7d6..874e9e5 100644 --- a/examples/flux_mnist/main.jl +++ b/examples/flux_mnist/flux_mnist.jl @@ -1,15 +1,15 @@ -using Flux, Zygote using Random, Statistics +using Flux using Flux.Losses: logitcrossentropy using Flux: onecold, onehotbatch using HuggingFaceDatasets -using MLUtils +using MLUtils: MLUtils, mapobs # using ProfileView, BenchmarkTools function mnist_transform(batch) - # under the "julia" (= numpy) format the image column is already a stacked - # (W, H, N) UInt8 array, so just rescale to Float32 in [0, 1] - image = Float32.(batch["image"]) ./ 255f0 + # the image column is a stacked (W, H, N) UInt8 array, + # so just rescale to Float32 in [0, 1] + image = batch["image"] ./ 255f0 return (; image, label = batch["label"]) end @@ -31,7 +31,7 @@ end # `num_workers = 0` loads on the main process; `num_workers > 0` spreads each batch's # `getobs` (and the CPython read it triggers) over that many worker processes, sidestepping # the GIL. MLUtils spawns the workers on demand under the current `--project`. -function train(; epochs=2, num_workers=0) +function train(; epochs=2, num_workers=0, materialize=false, parallel=false, verbose=true) batchsize = 128 nhidden = 100 device = cpu @@ -42,9 +42,13 @@ function train(; epochs=2, num_workers=0) # `num_workers > 0`); `mapobs`/`ObsView`-wrapped datasets compose with `num_workers` train_data = mapobs(mnist_transform, train_data) test_data = mapobs(mnist_transform, test_data) + if materialize + train_data = train_data[:] + test_data = test_data[:] + end - train_loader = Flux.DataLoader(train_data; batchsize, shuffle=true, num_workers) - test_loader = Flux.DataLoader(test_data; batchsize, num_workers) + train_loader = Flux.DataLoader(train_data; batchsize, shuffle=true, num_workers, parallel) + test_loader = Flux.DataLoader(test_data; batchsize, num_workers, parallel) model = Chain([Flux.flatten, Dense(28*28, nhidden, relu), @@ -61,17 +65,29 @@ function train(; epochs=2, num_workers=0) @info map(r, (; epoch, train_loss, train_acc, test_loss, test_acc)) end - report(0) - @time for epoch in 1:epochs + verbose && report(0) + for epoch in 1:epochs for (x, y) in train_loader x = x |> device yoh = onehotbatch(y, 0:9) |> device - loss, grads = withgradient(m -> logitcrossentropy(m(x), yoh), model) + loss, grads = Flux.withgradient(m -> logitcrossentropy(m(x), yoh), model) Flux.update!(opt, model, grads[1]) end - report(epoch) + verbose && report(epoch) end end -# @time train(; epochs=2, num_workers=0) -@time train(; epochs=2, num_workers=4) +println("#### START COMPARISON ###########") +MLUtils.close_dataloader_pool() +println("### WARMUP") # for precompilation +@time train(; epochs=1, num_workers=0, materialize=false, verbose=false) +println("### Serial") +@time train(; epochs=4, num_workers=0, materialize=false, verbose=false) +println("### Serial Materialized") +@time train(; epochs=4, num_workers=0, materialize=true, verbose=false) +println("### Paralle Materialized") +@time train(; epochs=4, num_workers=0, materialize=true, parallel=true, verbose=false) +println("### Distributed") +@time train(; epochs=4, num_workers=4, materialize=false, verbose=false) +MLUtils.close_dataloader_pool() +println("#### END COMPARISON ###########") diff --git a/examples/flux_mnist/pytorch_mnist.py b/examples/flux_mnist/pytorch_mnist.py new file mode 100644 index 0000000..9f4d910 --- /dev/null +++ b/examples/flux_mnist/pytorch_mnist.py @@ -0,0 +1,135 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["torch", "datasets", "numpy", "pillow"] +# /// +# +# PyTorch counterpart of flux_mnist.jl (Flux + HuggingFaceDatasets), for a data-loading timing +# comparison on the same task: an MLP on MNIST pulled from the same HuggingFace `datasets` +# Arrow dataset. Run with: uv run examples/flux_mnist/pytorch_mnist.py +# +# The four configs mirror flux_mnist.jl. One caveat: PyTorch's only DataLoader parallelism is +# multiprocess (`num_workers`), so the "Parallel Materialized" row uses worker *processes* +# over in-memory tensors — Julia's `parallel=true` there uses *threads* (cheap, shared +# memory), which PyTorch cannot do because of the GIL. This is a timing benchmark, not a +# numerical match: the exact AdamW weight decay etc. differ, so accuracies only track loosely. + +import time + +import numpy as np +import torch +import torch.nn as nn +from datasets import load_dataset, disable_progress_bars +from datasets.utils.logging import set_verbosity_error +from torch.utils.data import DataLoader, Dataset, TensorDataset + +disable_progress_bars() +set_verbosity_error() + +BATCHSIZE = 128 +NHIDDEN = 100 +EPOCHS = 4 +DEVICE = torch.device("cpu") + + +def make_model(): + return nn.Sequential( + nn.Flatten(), + nn.Linear(28 * 28, NHIDDEN), nn.ReLU(), + nn.Linear(NHIDDEN, NHIDDEN), nn.ReLU(), + nn.Linear(NHIDDEN, 10), + ).to(DEVICE) + + +class HFImageDataset(Dataset): + """On-the-fly loading: index the Arrow-backed HF dataset and decode per item. + + Mirrors `mapobs(mnist_transform, ds)` in flux_mnist.jl. The dataset is memory-mapped and + picklable, so with num_workers>0 each worker process re-opens it (data shared, not + copied) — the same mechanism as the Julia `num_workers` path. + """ + + def __init__(self, hf_ds): + self.ds = hf_ds.with_format("numpy") + + def __len__(self): + return len(self.ds) + + def __getitem__(self, i): + row = self.ds[i] + image = np.asarray(row["image"], dtype=np.float32) / 255.0 # (28, 28) in [0, 1] + return torch.from_numpy(image), int(row["label"]) + + +def materialize(hf_ds): + """Decode the whole split into in-memory tensors up front (the `[:]` path in flux_mnist.jl).""" + ds = hf_ds.with_format("numpy") + col = ds["image"] + images = col if isinstance(col, np.ndarray) else np.stack(col) + images = torch.from_numpy(images.astype(np.float32) / 255.0) # (N, 28, 28) + labels = torch.from_numpy(np.asarray(ds["label"]).astype(np.int64)) # (N,) + return TensorDataset(images, labels) + + +@torch.no_grad() +def loss_and_accuracy(loader, model): + model.eval() + lossfn = nn.CrossEntropyLoss(reduction="sum") + ls, correct, num = 0.0, 0, 0 + for x, y in loader: + x, y = x.to(DEVICE), y.to(DEVICE) + logits = model(x) + ls += lossfn(logits, y).item() + correct += (logits.argmax(1) == y).sum().item() + num += y.shape[0] + return ls / num, correct / num + + +def train(num_workers=0, materialize_data=False): + train_hf = load_dataset("ylecun/mnist", split="train") + test_hf = load_dataset("ylecun/mnist", split="test") + + if materialize_data: + train_ds, test_ds = materialize(train_hf), materialize(test_hf) + else: + train_ds, test_ds = HFImageDataset(train_hf), HFImageDataset(test_hf) + + # persistent_workers keeps the pool alive across epochs (closer to Julia's leased pool); + # without it PyTorch respawns workers on every pass over the loader. + kw = dict(num_workers=num_workers, persistent_workers=num_workers > 0) + train_loader = DataLoader(train_ds, batch_size=BATCHSIZE, shuffle=True, **kw) + test_loader = DataLoader(test_ds, batch_size=BATCHSIZE, **kw) + + model = make_model() + opt = torch.optim.AdamW(model.parameters(), lr=1e-3) + lossfn = nn.CrossEntropyLoss() + + def report(epoch): + tr_loss, tr_acc = loss_and_accuracy(train_loader, model) + te_loss, te_acc = loss_and_accuracy(test_loader, model) + print(f"(epoch = {epoch}, train_loss = {tr_loss:.3f}, train_acc = {tr_acc:.3f}, " + f"test_loss = {te_loss:.3f}, test_acc = {te_acc:.3f})") + + report(0) + for epoch in range(1, EPOCHS + 1): + model.train() + for x, y in train_loader: + x, y = x.to(DEVICE), y.to(DEVICE) + opt.zero_grad() + loss = lossfn(model(x), y) + loss.backward() + opt.step() + report(epoch) + + +def timed(name, **kw): + print(f"### {name}") + t0 = time.perf_counter() + train(**kw) + print(f" {time.perf_counter() - t0:.3f} seconds") + + +if __name__ == "__main__": + timed("Serial", num_workers=0, materialize_data=False) + timed("Serial Materialized", num_workers=0, materialize_data=True) + timed("Parallel Materialized", num_workers=4, materialize_data=True) + timed("Distributed", num_workers=4, materialize_data=False) diff --git a/examples/flux_mnist/pytorch_mnist_idiomatic.py b/examples/flux_mnist/pytorch_mnist_idiomatic.py new file mode 100644 index 0000000..ae8057b --- /dev/null +++ b/examples/flux_mnist/pytorch_mnist_idiomatic.py @@ -0,0 +1,151 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["torch", "torchvision", "datasets", "numpy", "pillow", "tqdm"] +# /// +# +# Idiomatic HuggingFace `datasets` + PyTorch MNIST training — written the way a practitioner +# typically would, rather than line-for-line equivalent to the Flux `flux_mnist.jl` (see +# `pytorch_mnist.py` for that). Run with: uv run examples/flux_mnist/pytorch_mnist_idiomatic.py +# +# Idioms shown: +# * `load_dataset` -> DatasetDict, then `ds["train"]` / `ds["test"]` +# * a torchvision transform applied lazily via `Dataset.with_transform` (the canonical image +# pipeline) — plus the `.map(batched=True)` + torch-format cached alternative +# * a `collate_fn` feeding a standard `DataLoader(num_workers=...)` +# * a plain `nn.Module`, `AdamW`, and a `tqdm` training loop +# +# The model, batch size and epochs match the Flux run so data-loading timings are comparable. +# Kept on CPU for that comparison; a typical script would auto-select the accelerator (see below). + +import time + +import numpy as np +import torch +import torch.nn as nn +from datasets import load_dataset, disable_progress_bars +from datasets.utils.logging import set_verbosity_error +from torch.utils.data import DataLoader +from torchvision.transforms import Compose, Normalize, ToTensor +from tqdm import tqdm + +disable_progress_bars() +set_verbosity_error() + +BATCHSIZE = 128 +EPOCHS = 4 +# A typical script auto-selects the accelerator: +# DEVICE = torch.device("cuda" if torch.cuda.is_available() +# else "mps" if torch.backends.mps.is_available() else "cpu") +# Pinned to CPU here so the timings line up with the Flux example. +DEVICE = torch.device("cpu") + +# ToTensor scales uint8 [0, 255] -> float [0, 1] and lays out (C, H, W); Normalize applies the +# standard MNIST mean/std. This is the usual torchvision image pipeline. +IMAGE_TF = Compose([ToTensor(), Normalize((0.1307,), (0.3081,))]) + + +class MLP(nn.Module): + def __init__(self, nhidden=100): + super().__init__() + self.net = nn.Sequential( + nn.Flatten(), + nn.Linear(28 * 28, nhidden), nn.ReLU(), + nn.Linear(nhidden, nhidden), nn.ReLU(), + nn.Linear(nhidden, 10), + ) + + def forward(self, x): + return self.net(x) + + +def lazy_transform(batch): + """Decode + transform on access (registered with `with_transform`).""" + batch["pixel_values"] = [IMAGE_TF(img) for img in batch["image"]] + return batch + + +def collate(examples): + x = torch.stack([e["pixel_values"] for e in examples]) + y = torch.tensor([e["label"] for e in examples]) + return x, y + + +def cached_preprocess(batch): + """Same transform, but returning arrays so `datasets` can cache them to Arrow.""" + batch["pixel_values"] = [IMAGE_TF(img).numpy() for img in batch["image"]] + return batch + + +def make_loaders(num_workers, cached): + ds = load_dataset("ylecun/mnist") + if cached: + # Preprocess once; `datasets` caches the result to Arrow and hands back torch tensors + # on access, so the default collate works and no per-item Python decode is needed. + ds = ds.map(cached_preprocess, batched=True, remove_columns=["image"]) + ds = ds.with_format("torch") + train_split, test_split, collate_fn = ds["train"], ds["test"], None + else: + train_split = ds["train"].with_transform(lazy_transform) + test_split = ds["test"].with_transform(lazy_transform) + collate_fn = collate + + kw = dict(num_workers=num_workers, persistent_workers=num_workers > 0, collate_fn=collate_fn) + train_loader = DataLoader(train_split, batch_size=BATCHSIZE, shuffle=True, **kw) + test_loader = DataLoader(test_split, batch_size=BATCHSIZE, **kw) + return train_loader, test_loader + + +def _xy(batch): + # lazy path yields (x, y) tuples; cached torch-format path yields column dicts. + if isinstance(batch, dict): + return batch["pixel_values"], batch["label"] + return batch + + +@torch.no_grad() +def evaluate(loader, model, lossfn): + model.eval() + total_loss, correct, num = 0.0, 0, 0 + for batch in loader: + x, y = _xy(batch) + x, y = x.to(DEVICE), y.to(DEVICE) + logits = model(x) + total_loss += lossfn(logits, y).item() * y.shape[0] + correct += (logits.argmax(1) == y).sum().item() + num += y.shape[0] + return total_loss / num, correct / num + + +def train(num_workers=0, cached=False): + train_loader, test_loader = make_loaders(num_workers, cached) + model = MLP().to(DEVICE) + opt = torch.optim.AdamW(model.parameters(), lr=1e-3) + lossfn = nn.CrossEntropyLoss() + + for epoch in range(1, EPOCHS + 1): + model.train() + for batch in tqdm(train_loader, desc=f"epoch {epoch}/{EPOCHS}", leave=False): + x, y = _xy(batch) + x, y = x.to(DEVICE), y.to(DEVICE) + opt.zero_grad() + loss = lossfn(model(x), y) + loss.backward() + opt.step() + tr_loss, tr_acc = evaluate(train_loader, model, lossfn) + te_loss, te_acc = evaluate(test_loader, model, lossfn) + print(f"(epoch = {epoch}, train_loss = {tr_loss:.3f}, train_acc = {tr_acc:.3f}, " + f"test_loss = {te_loss:.3f}, test_acc = {te_acc:.3f})") + + +def timed(name, **kw): + print(f"### {name}") + t0 = time.perf_counter() + train(**kw) + print(f" {time.perf_counter() - t0:.3f} seconds") + + +if __name__ == "__main__": + timed("Lazy with_transform, num_workers=0", cached=False, num_workers=0) + timed("Lazy with_transform, num_workers=4", cached=False, num_workers=4) + timed("Cached .map (torch format), num_workers=0", cached=True, num_workers=0) + timed("Cached .map (torch format), num_workers=4", cached=True, num_workers=4) From a8c60cb43b1a57ae15b318992db084af2f642023 Mon Sep 17 00:00:00 2001 From: CarloLucibello Date: Mon, 6 Jul 2026 10:04:51 +0200 Subject: [PATCH 7/8] CHANGELOG: split Unreleased into 0.4.0, 0.4.1, and Unreleased MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 0.4.1 gets the Features/ClassLabel schema views (#65) and MLCore ≥ 1.1 getobs(::Py) change (#66); 0.4.0 keeps the breaking-release content; the num_workers DataLoader entry stays under Unreleased. Compare links updated. Co-Authored-By: Claude Opus 4.8 (1M context) --- CHANGELOG.md | 53 +++++++++++++++++++++++++++++----------------------- 1 file changed, 30 insertions(+), 23 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4c9da59..37d612c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,11 +9,34 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - `MLUtils.DataLoader(ds::Dataset; num_workers=N)` works out of the box for process-parallel - loading, including over `mapobs`/`ObsView`-wrapped datasets. The GIL-safe serialization it - relies on — running the dataset's `pickle` on the main task and loading this package on the - workers — is handled upstream by MLUtils, so no wrapper type is exported here. An `@info` notes - when an in-memory dataset is first materialized to a temporary Arrow file for this. Requires - MLUtils ≥ 0.4.11. + loading, including over `mapobs`/`ObsView`-wrapped datasets. Building on the `Serialization` + support added in 0.4.0, the GIL-safe serialization it relies on — running the dataset's + `pickle` on the main task and loading this package on the workers — is handled upstream by + MLUtils, so no wrapper type is exported here. An `@info` notes when an in-memory dataset is + first materialized to a temporary Arrow file for this. Requires MLUtils ≥ 0.4.11. + +## [0.4.1] - 2026-07-06 + +### Added +- Julia views over a dataset's schema. `ds.features` now returns a `Features` view (an + `AbstractDict{String, Any}`) instead of a raw `Py` mapping: indexing a column yields a wrapped + `ClassLabel`/`Value` leaf (other feature types stay raw `Py`), each forwarding attribute and + method access to Python (`cl.names`, `cl.num_classes`, `cl.int2str(i)`, `cl.str2int(s)`, + `v.dtype`). The views can be built from Julia (`ClassLabel(names=[…])`, `Value("int64")`, + `Features(Dict(…))`) and handed back to Python via `jl2py` (e.g. a `features=` schema + argument). Public but unexported Julian conveniences look up a column's `ClassLabel` in one + call: `class_names(ds, col)`, `int2str(ds, col, i)`, `str2int(ds, col, s)` (label ids are + 0-based data, passed through with no index offset), plus `features(ds)` as the function form + of `ds.features`. `Features`/`ClassLabel`/`Value` and these helpers are all public but not + exported — the Pythonic idioms (`ds.features`, method chaining, `datasets.ClassLabel(…)`) are + the primary interface. + +### Changed +- Depend on `MLCore` (≥ 1.1) instead of `MLUtils` for the `getobs`/`numobs` observation + interface, and drop the package's own `getobs(::Py, ::Integer)` method (type piracy on + `PythonCall.Py`). `getobs` for Python observation containers is now provided by MLCore's + `PythonCall` extension. `MLUtils.DataLoader` still works, since MLUtils re-exports the + same `getobs`/`numobs` from MLCore. ## [0.4.0] - 2026-07-02 @@ -47,18 +70,6 @@ Julia values instead of raw Python objects. See **Breaking** below before upgrad `BoundsError`, instead of `AssertionError` — update any code catching `AssertionError`. ### Added -- Julia views over a dataset's schema. `ds.features` now returns a `Features` view (an - `AbstractDict{String, Any}`) instead of a raw `Py` mapping: indexing a column yields a wrapped - `ClassLabel`/`Value` leaf (other feature types stay raw `Py`), each forwarding attribute and - method access to Python (`cl.names`, `cl.num_classes`, `cl.int2str(i)`, `cl.str2int(s)`, - `v.dtype`). The views can be built from Julia (`ClassLabel(names=[…])`, `Value("int64")`, - `Features(Dict(…))`) and handed back to Python via `jl2py` (e.g. a `features=` schema - argument). Public but unexported Julian conveniences look up a column's `ClassLabel` in one - call: `class_names(ds, col)`, `int2str(ds, col, i)`, `str2int(ds, col, s)` (label ids are - 0-based data, passed through with no index offset), plus `features(ds)` as the function form - of `ds.features`. `Features`/`ClassLabel`/`Value` and these helpers are all public but not - exported — the Pythonic idioms (`ds.features`, method chaining, `datasets.ClassLabel(…)`) are - the primary interface. - `Serialization` support for `Dataset`: `Serialization.serialize`/`deserialize` now work, so a `Dataset` can cross a process boundary — the prerequisite for process-parallel data loaders (e.g. a `MLUtils.DataLoader(ds; num_workers=N)` that spreads `getobs` over worker @@ -163,11 +174,6 @@ Julia values instead of raw Python objects. See **Breaking** below before upgrad `Dataset` (was `pyds`) and `DatasetDict` (was `pyd`), for consistency. This field is not part of the public API (it is shadowed by `getproperty`), so the change is non-breaking for documented usage. -- Depend on `MLCore` (≥ 1.1) instead of `MLUtils` for the `getobs`/`numobs` observation - interface, and drop the package's own `getobs(::Py, ::Integer)` method (type piracy on - `PythonCall.Py`). `getobs` for Python observation containers is now provided by MLCore's - `PythonCall` extension. `MLUtils.DataLoader` still works, since MLUtils re-exports the - same `getobs`/`numobs` from MLCore. ### Fixed - Fixed a segfault triggered by REPL tab-completion on a `DatasetDict` (`d[`), which @@ -198,6 +204,7 @@ Baseline. Changes up to and including this release are recorded in the [git history](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/commits/main) and the [GitHub releases](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/releases). -[Unreleased]: https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/compare/v0.4.0...HEAD +[Unreleased]: https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/compare/v0.4.1...HEAD +[0.4.1]: https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/compare/v0.4.0...v0.4.1 [0.4.0]: https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/compare/v0.3.4...v0.4.0 [0.3.4]: https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/releases/tag/v0.3.4 From 36e7da58ee2a5c06ba035bf4e7cdeab4a93abe5c Mon Sep 17 00:00:00 2001 From: CarloLucibello Date: Mon, 6 Jul 2026 10:04:51 +0200 Subject: [PATCH 8/8] docs: expand Data Loaders guide with runnable examples 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) --- docs/src/guide.md | 87 +++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 80 insertions(+), 7 deletions(-) diff --git a/docs/src/guide.md b/docs/src/guide.md index c38cf35..fc152d0 100644 --- a/docs/src/guide.md +++ b/docs/src/guide.md @@ -506,23 +506,96 @@ A [`Dataset`](@ref) is a valid MLUtils data container, so you can hand one strai [MLUtils Data Loaders guide](https://juliaml.github.io/MLUtils.jl/stable/guide/iteration/); the notes below are specific to a `Dataset`. +### Iterating batches + +Each batch is the dataset's own `"julia"`-format batch: the columns for the batch's rows, +stacked and keyed by name (numeric array columns stack into an N-D array with the +observation axis last). Small, self-contained example: + +```jldoctest guide +julia> loader = DataLoader(Dataset((; x=1:6, y=10:10:60)); batchsize=2); + +julia> batch = first(loader); + +julia> batch["x"], batch["y"] +([1, 2], [10, 20]) +``` + +Most training loops want `(input, target)` pairs rather than a `Dict`, so wrap the dataset +in a `mapobs` that reshapes each observation; the loader then collates the tuples +column-wise into an `(inputs, targets)` batch you can destructure with `for (x, y) in +loader`: + +```jldoctest guide +julia> reshape_obs(o) = (o["x"] * 100, o["y"]); # per-observation: Dict -> (input, target) + +julia> loader = DataLoader(mapobs(reshape_obs, Dataset((; x=1:6, y=10:10:60))); batchsize=2); + +julia> first(loader) +([100, 200], [10, 20]) +``` + +Against the Hub, the same shapes hold — swap the in-memory `Dataset` for a `load_dataset`: + ```julia-repl julia> using MLUtils julia> ds = load_dataset("ylecun/mnist", split="train"); # "julia" format by default julia> loader = DataLoader(ds; batchsize=128, shuffle=true); + +julia> batch = first(loader); # batch["image"] is (28, 28, 128), batch["label"] is length-128 ``` -For small datasets you can **materialize** into memory up front with `[:]` (faster per epoch, -no decoding during iteration); otherwise the loader reads **on the fly**, keeping memory flat. +### Materializing vs. reading on the fly + +For small datasets you can **materialize** into memory up front with `[:]` (faster per +epoch, no decoding during iteration); otherwise the loader reads **on the fly**, keeping +memory flat: + +```julia-repl +julia> materialized = DataLoader(ds[:]; batchsize=128, shuffle=true); # whole dataset in RAM + +julia> on_the_fly = DataLoader(ds; batchsize=128, shuffle=true); # decode per batch +``` + +Layer a per-batch transform with `mapobs`; with `[:]` on the *mapped* dataset the transform +runs once up front, without it the transform runs per batch during iteration: + +```julia-repl +julia> mnist_transform(o) = (Float32.(o["image"]) ./ 255, o["label"]); + +julia> train = mapobs(mnist_transform, ds); + +julia> loader = DataLoader(train[:]; batchsize=128, shuffle=true); # or drop [:] for on-the-fly +``` + +### Parallel loading (`num_workers`) Reading on the fly calls into CPython, and CPython's global interpreter lock (**GIL**) serializes those reads, so thread-based `parallel=true` gives ~1× on a `Dataset`. Use `num_workers=N` instead (needs MLUtils ≥ 0.4.11): it spreads `getobs` over `N` worker -**processes**, each with its own interpreter and GIL, so loading scales past the GIL. This -works because a `Dataset` serializes by *reference* to its on-disk Arrow files (an in-memory -dataset is first materialized to a temporary Arrow file), so it needs a serializable `jltransform` if you set one. See the MLUtils guide for the mechanics and tradeoffs. +**processes**, each with its own interpreter and GIL, so loading scales past the GIL: + +```julia-repl +julia> loader = DataLoader(ds; batchsize=128, num_workers=4); # 4 worker processes + +julia> for batch in loader + # ... training step ... + end +``` -A complete example — `mapobs` transform, `num_workers`, and a training loop — lives in -[`examples/flux_mnist/flux_mnist.jl`](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/blob/main/examples/flux_mnist/flux_mnist.jl). +This works because a `Dataset` serializes by *reference* to its on-disk Arrow files +(an in-memory dataset is first materialized to a temporary Arrow file, announced with an +`@info`), so if you set a `jltransform` it must be serializable. It composes with a +`mapobs`/`ObsView` wrapper at any nesting — `DataLoader(mapobs(f, ds); num_workers=4)` runs +`f` on the workers. Batches may arrive **out of input order** across workers; combine with +`shuffle=true` (or don't rely on order). Worker processes are launched on demand under the +current `--project`; call `MLUtils.close_dataloader_pool()` to tear them down. Process +parallelism only pays off when `getobs` is the bottleneck (large images, heavy decode) — +for cheap observations the IPC overhead can make it slower than serial. See the MLUtils +guide for the mechanics and tradeoffs. + +A complete example — `mapobs` transform, `num_workers`, and a training loop, benchmarked +against PyTorch — lives in +[`examples/flux_mnist/`](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/tree/main/examples/flux_mnist).