The PyO3 module dataloader_rs, built by maturin into the dataloader-rs Python package. This is
the implementation guide; the user guide is python/README.md.
In this guide, a thread attaches to the interpreter when it takes the right to run Python code
(the GIL, on builds that have one) and detaches when it gives it back; these are PyO3's terms
(Python::attach, py.detach).
| File | What |
|---|---|
module.rs |
The #[pymodule], declared with gil_used = false so that it loads on free-threaded CPython without enabling the GIL. |
error.rs |
into_py_err, which turns the core's errors back into Python exceptions. |
dataset/sync.rs |
PyDatasetBase (the PyDataset base class), Dataset for Py<PyDatasetBase>, Fetcher (__getitems__ or __getitem__) |
dataloader.rs |
PyDataloader: arguments, validation, __iter__, __len__ |
iterator.rs |
PyDataloaderIter: the direct, threaded and async epochs, __next__ |
transform.rs |
PyTransform: transform=f, called as f(index, item), or the Rust image transforms (Native); the loader's Transform for Python items and for a Rust dataset's Bytes; PyItem, a Python object or a Rust value |
sampler.rs |
PySampler: sequential, random, or a Python iterable, and the exception it raised |
collator.rs |
PyCollator: collate_fn, an ImageCollator (collated in Rust), or a plain list |
image/ops.rs |
Value (bytes or an image of one type and layout) and Op (one Rust image transform), applied without static types |
image/transforms.rs |
dataloader_rs.transforms: Transform (a list of Ops, callable) and its constructors Decode, PadTo, …, Compose |
image/collator.rs |
ImageCollator in Python: Rust values, or Image objects and (Image, …) tuples |
image/batch.rs |
Image and ImageBatch: __array_interface__ views and __dlpack__ |
image/array.rs |
Image.from_array: a numpy array or CPU tensor (through np.asarray and the buffer protocol), checked against Decode's limits, copied into an HWC image |
pool.rs |
run_attached: the pool threads' main function, one thread state per thread; attach_global_pool does the same for rayon's global pool |
dataset/asynchronous.rs |
PyAsyncDatasetBase (the PyAsyncDataset base class), PyAsyncItems: coroutines on an asyncio loop thread, submitted from the global pool |
dataset/event_loop.py |
that loop: start(aclose) returns submit(coroutine) and stop(); compiled into the module |
dataset/rust.rs |
RustDataset and its HttpDataset / FileDataset / ObjectStoreDataset subclasses |
dataset/kind.rs |
a dataset's kind from its base class, and the checks its methods must pass |
iter(loader)
├─ len(dataset) raises if __len__ raises
├─ sampler → batches of indices a Python sampler is iterated here, under the interpreter
├─ num_workers == 0, item_concurrency == 1 → Direct __next__ fetches, transforms, collates
├─ otherwise → Threaded the Rust core; one attach per batch, or per item
└─ PyAsyncDataset or Rust dataset → Async the async core (futures on tokio), always
- Direct: the
Fetcher(a bound__getitems__, or a bound__getitem__) is resolved once per epoch;__next__calls it, thentransformper item andcollate_fn. - Threaded: the core's worker threads call
Dataset::get_batchonPy<PyDatasetBase>, which attaches to the interpreter once for the whole batch and resolves theFetcherthere.__next__returns a batch that already arrived directly, and releases the interpreter (py.detach) only when it has to wait on the channel, in 100 ms slices that end withcheck_signals(Ctrl-C). - Pool threads keep a thread state (
pool.rs): each attaches once when it starts and runs rayon's loop detached, so attaching from a task re-activates its thread state instead of creating and destroying one, which costs about 100 µs on free-threaded builds. item_concurrency > 1: the epoch uses the core iterator even withnum_workers = 0; the batch's lanes fetch its items with oneDataset::get(one attach, one__getitem__) each.- Async:
PyDataloaderbuildsDataLoader::builder_asyncaroundPyAsyncItemsor a Rust dataset, withPyTransform(which also makes a Rust dataset's itemsbytes);__next__waits on it like on the threaded core (wait_next). See ARCHITECTURE.md, 6.3 and 6.4.
- Never wait for workers while holding the interpreter. A worker may be waiting to attach
(inside
__getitem__orcollate_fn);__next__detaches before it waits on the channel, and dropping an iterator detaches before it waits (at most 5 s) for the batches in flight. - Python exceptions keep their type. They cross the core as boxed
PyErrs;into_py_errdowncasts them back, and only non-Python errors becomeRuntimeError. - Nothing ends an epoch silently. A raising
__len__, a sampler that raises or yields a non-index, a sampler without__len__inlen(loader): each raises its own exception instead of producing an empty or shorter epoch. - The loader outlives its iterators.
PyDataloaderIterholds a strong reference to itsPyDataloader. - The garbage collector sees every Python reference, once.
__traverse__reports the dataset, the Python sampler,transformandcollate_fnfrom the loader, and_ownerplus the direct path's own references from an iterator, so cycles through them are freed. Running epochs share the loader's objects throughArcs and hold no references of their own; reporting one twice would let the collector free a live object.