datago

A Rust-written data loader which can be used as a python module. Handles several data sources, from local files to webdataset or a VectorDB focused http stack like our open sourced dataroom. Focused on image data at the moment, could also easily be more generic.

Datago handles, outside of the Python GIL

• per sample IO
• deserialization (jpg and png decompression)
• some optional vision processing (aligning different image payloads)
• optional serialization

Samples are exposed in the Python scope as python native objects, using PIL and Numpy base types. Speed will be network dependent, but GB/s is typical. Depending on the front ends, datago can be rank and world-size aware, in which case the samples are dispatched depending on the samples hash.

Use it

You can simply install datago with [uv] pip install datago

Use the package from Python

Please note that in all the of the following cases, you can directly get an IterableDataset (torch compatible) with the following code snippet

from dataset import DatagoIterDataset

client_config = {}  # See below for examples
datago_dataset = DatagoIterDataset(client_config, return_python_types=True)

return_python_types enforces that images will be of the PIL.Image sort for instance, being an external binary module should be transparent.

Python sample reads try the ready queue without releasing the GIL, then release it while waiting and poll Python signals every 100 ms. Python delivers SIGINT handlers on the main thread only. If a reader runs in a background thread, handle Ctrl-C on the main thread by calling client.stop(); that cancellation is safe while the reader is waiting. stop() closes both pipeline queues, releases the GIL while joining, and returns after a short grace period if native work is still in flight. Remaining thread handles are retained by a cleanup reaper rather than detached. Prefer try/finally with client.stop() for explicit cleanup. Queued async work is cancelled; a filesystem/network syscall already executing cannot be forcibly stopped and may continue in the background until the OS call returns. The native Rust client retains its blocking API. Reads do not silently restart: once the stream is exhausted (or stop() has been called) get_sample() returns None, and you must call start() again to begin a new pass.

Waiting on an empty queue also lets other Python threads run. For very cheap samples, GIL contention can reduce throughput compared with a reader that keeps the GIL; queued samples still use the fast path without releasing it. live_cleanup_threads() counts cleanup reapers, not all native threads or blocking-pool tasks that outlive runtime shutdown.

Read results and errors

Outcome	Rust (try_get_sample)	Python (get_sample / get_sample_auto_convert)
Sample available	Ok(Some(sample))	Return the sample
Normal exhaustion or explicit stop	Ok(None)	Return None
Fatal source or worker failure	Err(DatagoError::Source(_)) / Err(DatagoError::Worker(_))	Raise RuntimeError
Read timeout	Err(DatagoError::Timeout)	Raise TimeoutError
Invalid configuration or timeout argument	Err(DatagoError::Configuration(_))	Raise ValueError

Unreadable individual images (including decoder panics) and unreadable nested directories are logged and skipped. An unreadable filesystem root, failed metadata request/parse, or worker failure ends the pass with an error. WDS listing errors and shard download/unpacking errors after retries are fatal: Datago does not silently skip a missing or corrupt shard and drop its samples. A fatal failure cancels outstanding work and discards queued samples. Concurrent stop and failure use the first recorded terminal outcome. After a fatal source/worker error, subsequent reads return None; call start() explicitly to begin a new pass. Errors belong to their engine generation and cannot stop a restarted pass.

Python reads accept a keyword-only timeout in seconds (default: 300). Rust callers can use try_get_sample_timeout(Duration) for a per-read timeout. A timeout only ends that read, including at the default 300-second deadline: retry to keep reading the same pass, or call stop() to cancel it. Zero returns a ready sample or raises TimeoutError immediately if the queue is empty and still open; an exhausted or stopped stream returns None. Cleanup gets a short grace period; stalled native operations may finish under a background reaper. Python KeyboardInterrupt remains an interrupt rather than a pipeline error.

try:
    while (sample := client.get_sample(timeout=30.0)) is not None:
        process(sample)
finally:
    client.stop()

The existing Rust get_sample() remains a compatibility wrapper: it logs failures and returns None. Use try_get_sample() when failure must be distinguished from normal exhaustion. Python propagates failures through its existing read methods.

Release notes: 2026.10.2

• Breaking change: Python get_sample() and get_sample_auto_convert() now raise RuntimeError for fatal source/worker failures instead of returning None. This includes DB outages and WDS shard failures after retries.
• Python start() and reads raise ValueError for invalid client/source configuration; invalid read timeouts also raise ValueError.
• Reads raise TimeoutError on their deadline (300 seconds by default). The pass stays active: retry the read to continue, or call stop() to cancel it.
• Rust callers can use try_get_sample() / try_get_sample_timeout() for typed errors. The legacy Rust get_sample() still logs errors and returns None.
• Low-level Rust source orchestrate() functions now return Result<DatagoEngine, DatagoError>. Use engine.close() or drop the engine to cancel it; its metadata queue and cancellation token are now internal.

Release notes: 2026.10.1

• Breaking change: reads after end-of-stream or stop() return None rather than silently starting another epoch. Call start() explicitly for each new pass.
• DB processing now defaults to min(CPU count, connection budget, sample limit) in-flight tasks, rather than accidentally serializing when DATAGO_MAX_TASKS is unset. WDS processing defaults to one task per CPU. DATAGO_MAX_TASKS=0 retains serial processing; explicit overrides retain their existing source-specific caps.
• Python waits release the GIL and check signals; cancellation also covers DB page requests and WDS shard listing/downloads. Already-running blocking filesystem calls may still outlive shutdown.
• Source builds require maturin>=1.9.4,<2. Release builds no longer use panic = "abort", allowing worker panics to be reported instead of aborting the process.

Dataroom

from datago import DatagoClient, initialize_logging
import os
import json

# Respects RUST_LOG=INFO env var for setting log level
# If omitted the logger will be initialized when the client starts.
initialize_logging()

config = {
    "source_config": {
        "sources": os.environ.get("DATAROOM_TEST_SOURCE", ""),
        "page_size": 500,
        "rank": 0,
        "world_size": 1,
    },
    "limit": 200,
    "samples_buffer_size": 32,
}

client = DatagoClient(json.dumps(config))

for _ in range(10):
    sample = client.get_sample()

Please note that the image buffers will be passed around as raw pointers, see below (we provide python utils to convert to PIL types).

Local files

To test datago while serving local files (jpg, png, ..), code would look like the following. Note that datago serving files with a lot of concurrent threads means that, even if random_sampling is not set, there will be some randomness in the sample ordering.

from datago import DatagoClient, initialize_logging
import os
import json

# Can also set the log level directly instead of using RUST_LOG env var
initialize_logging(log_level="warn")

config = {
    "source_type": "file",
    "source_config": {
        "root_path": "myPath",
        "random_sampling": False,  # True if used directly for training
        "rank": 0,  # Optional, distributed workloads are possible
        "world_size": 1,
    },
    "limit": 200,
    "samples_buffer_size": 32,
}

client = DatagoClient(json.dumps(config))

for _ in range(10):
    sample = client.get_sample()

[experimental] Webdataset

Please note that this implementation is very new, and probably has significant limitations still. It has not yet been tested at scale. Please also note that you can find a better example in /python/benchmark_webdataset.py, which will show how to convert everything to more pythonic types (PIL images).

from datago import DatagoClient, initialize_logging
import os
import json

# Can also set the log level directly instead of using RUST_LOG env var
initialize_logging(log_level="warn")

# URL of the test bucket
bucket = "https://storage.googleapis.com/webdataset/fake-imagenet"
dataset = "/imagenet-train-{000000..001281}.tar"
url = bucket + dataset

client_config = {
    "source_type": "webdataset",
    "source_config": {
        "url": url,
        "random_sampling": False,
        "concurrent_downloads": 8,  # The number of TarballSamples which should be handled concurrently
        "rank": 0,
        "world_size": 1,
    },
    "prefetch_buffer_size": 128,
    "samples_buffer_size": 64,
    "limit": 1_000_000,  # Dummy example, max number of samples you would like to serve
}

client = DatagoClient(json.dumps(client_config))

for _ in range(10):
    sample = client.get_sample()

Process images on the fly

Datago can also process images on the fly, for instance to align different image payloads. This is done by adding an image_config to the configuration. The following example shows how to align different image payloads.

Processing can be very CPU heavy, but it will be distributed over all CPU cores wihout requiring multiple python processes. I.e., you can keep a single python process using get_sample() on the client and still saturate all CPU cores.

There are three main processing topics that you can choose from:

• crop the images to within an aspect ratio bucket (which is very handy for all Transformer / patch based architectures)
• resize the images (setting here will be related to the square aspect ratio bucket, other buckets will differ of course)
• pre-encode the images to a specific format (jpg, png, ...)

   config = {
    "source_type": "file",
    "source_config": {
        "root_path": "myPath",
        "random_sampling": False, # True if used directly for training
    },
    # Optional pre-processing of the images, placing them in an aspect ratio bucket to preserve as much as possible of the original content
    "image_config": {
        "crop_and_resize": True, # False to turn it off, or just omit this part of the config
        "default_image_size": 1024,
        "downsampling_ratio": 32,
        "min_aspect_ratio": 0.5,
        "max_aspect_ratio": 2.0,
        "pre_encode_images": False,
    },
    "limit": 200,
    "samples_buffer_size": 32,
}

Match the raw exported buffers with typical python types

See helper functions provided in raw_types.py, should be self explanatory. Check python benchmarks for examples. As mentioned above, we also provide a wrapper so that you get a dataset directly.

Logging

We are using the log crate with env_logger. You can set the log level using the RUST_LOG environment variable. E.g. RUST_LOG=INFO.

When using the library from Python, env_logger will be initialized automatically when creating a DatagoClient. There is also a initialize_logging function in the datago module, which if called before using a client, allows to customize the log level. This only works if RUST_LOG is not set.

Env variables

There are a couple of env variables which will change the behavior of the library, for settings which felt too low level to be exposed in the config.

• DATAGO_MAX_TASKS: refers to the number of threads which will be used to load the samples. Defaults to a multiple of the CPU cores.
• RUST_LOG: see above, will change the level of logging for the whole library, could be useful for debugging or to report an issue here.
• DATAGO_MAX_RETRIES: number of retries for a failed sample load, defaults to 3.

Build it

Preamble

Just install the rust toolchain via rustup

[Apple Silicon MacOS only]

If you are using an Apple Silicon Mac OS machine, create a .cargo/config.toml file and paste the following:

[target.x86_64-apple-darwin]
rustflags = [
  "-C", "link-arg=-undefined",
  "-C", "link-arg=dynamic_lookup",
]

[target.aarch64-apple-darwin]
rustflags = [
  "-C", "link-arg=-undefined",
  "-C", "link-arg=dynamic_lookup",
]

Build a benchmark CLI

cargo run --release -- -h to get all the information, should be fairly straightforward

Native filesystem benchmark

Run the filesystem pipeline directly from Rust, including image decoding and optional preprocessing, without Python conversion:

cargo run --release --example benchmark_filesystem -- /path/to/images --limit 10000 --num-workers 16

Add --crop-and-resize for a 1024 × 1024 center crop. --num-workers controls in-flight tasks and defaults to CPU count; 0 means serial processing. The benchmark uses a 256-sample buffer and requires the requested number of samples. Image payloads are borrowed when collecting metrics, avoiding a pixel-buffer copy.

The JSON report includes FPS, output MP/s and MB/s, time to the first sample, and FPS after the first sample (steady_fps, or null for one sample). Elapsed time includes construction and startup and excludes cleanup. Compare repeated release runs with identical options and comparable filesystem cache state.

Run the rust test suite

From the datago folder

cargo test

Generate the python package binaries manually

Build a wheel useable locally

maturin build -i python3.11 --release --target "x86_64-unknown-linux-gnu"

Build a wheel which can be uploaded to pypi or related

• either use a manylinux docker image

• or cross compile using zip

maturin build -i python3.11 --release --target "x86_64-unknown-linux-gnu" --manylinux 2014 --zig

then you can pip install from target/wheels

Update the pypi release (maintainers)

Create a new tag and a new release in this repo, a new package will be pushed automatically.

Benchmarks

As usual, benchmarks are a tricky game, and you shouldn't read too much into the following plots but do your own tests. Some python benchmark examples are provided in the [python](./python/) folder.

In general, Datago will be impactful if you want to load a lot of images very fast, but if you consume them as you go at a more leisury pace then it's not really needed. The more CPU work there is with the images and the higher quality they are, the more Datago will shine.

From disk: ImageNet

The following benchmarks are using ImageNet 1k, which is very low resolution and thus kind of a worst case scenario. Data is served from cache (i.e. the OS cache) and the images are not pre-processed. In this case the receiving python process is typically the bottleneck, and caps at around 3000 images per second.

AMD Zen3 laptop - IN1k - disk - no processing

AMD EPYC 9454 - IN1k - disk - no processing

One liner to repro locally (rebuilding the package):

DATAGO_TEST_FILESYSTEM=$PATH_TO_YOUR_TEST_DATA uv run --reinstall-package datago --python 3.14 --group dev python/benchmark_filesystem.py --sweep --limit 500

Webdataset: FakeIN

This benchmark is using low resolution images. It's accessed through the webdataset front end, datago is compared with the popular python webdataset library. Note that datago will start streaming the images faster here (almost instantly !), which emphasizes throughput differences depending on how long you test it for.

Of note is also that this can be bottlenecked by your external bandwidth to the remote storage where WDS is hosted, in which case both solution would yield comparable numbers.

AMD Zen3 laptop - webdataset - no processing

Webdataset: PD12M

This benchmark is using high resolution images. It's accessed through the webdataset front end, datago is compared with the popular python webdataset library. Note that datago will start streaming the images faster here (almost instantly !), which emphasizes throughput differences depending on how long you test it for.

Of note is also that this can be bottlenecked by your external bandwidth to the remote storage where WDS is hosted, in which case both solution would yield comparable numbers.

AMD Zen3 laptop - webdataset - no processing

AMD EPYC 9454 - pd12m - webdataset - no processing

AMD Zen3 laptop - webdataset - processing

Adding image processing (crop and resize to Transformer compatible size buckets) to the equation changes the picture, as the work spread becomes more important. If you're training a diffusion model or an image encoder from a diverse set of images, this is likely to be the most realistic micro-benchmark.

License

MIT License

Copyright (c) 2025 Photoroom

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.