Multiprocessing and Multithreading

AudioTree datasets are Grain datasets, so Grain’s two parallelism mechanisms, multiprocessing via mp_prefetch and multithreading via ReadOptions, apply to them unchanged. This page is about when each one actually helps.

Note

Multiprocessing is Linux and macOS only. Grain’s workers return results through named shared memory, and Windows destroys a named mapping as soon as its last handle closes – so the parent fails to attach with FileNotFoundError: [WinError 2] ... 'wnsm_<id>' once the worker exits. This is google/grain#793, still open upstream, and not something AudioTree can work around.

Multithreading via ReadOptions(num_threads=...) is unaffected – threads share one address space and need no shared memory. On Windows, raise num_threads instead of adding workers, or run under WSL.

Why Parallel Loading?

Loading audio is real work: reading files off disk, decoding FLAC or MP3, resampling, and (with saliency enabled) measuring loudness. Done serially on the training thread, that work leaves the GPU waiting between steps; done in parallel workers, it overlaps with training.

Synchronous reads are the right baseline

Grain’s default ReadOptions is num_threads=16, prefetch_buffer_size=500 – sized for corpora of small records, not decoded audio. A 3-second stereo excerpt at 44.1 kHz float32 is ~1 MB, so the default buffer alone can hold ~0.5 GB of decoded audio, and it applies per worker process once mp_prefetch is involved (8 workers -> ~4 GB of prefetched waveforms). The read threads also compete with the training process for the GIL.

Start every pipeline with fully synchronous reads, and add parallelism only where measurement says it pays:

read_options = grain.ReadOptions(num_threads=0, prefetch_buffer_size=0)
iter_ds = ds.to_iter_dataset(read_options=read_options)

• Training with mp_prefetch: keep reads synchronous inside each worker – the worker processes are the parallelism, and per-worker thread pools mostly add memory and GIL churn.

• Evaluation and pre-rendering scripts: synchronous reads keep memory flat and behavior deterministic.

• Genuinely I/O-bound reads (network storage, cold disks) are the case for num_threads > 0 – the next section – and even then prefer a small explicit prefetch_buffer_size over the default 500.

Multithreading with ReadOptions

Use multithreading for I/O-bound operations:

import grain
from audiotree.sources import create_balanced_audio_dataset

# Create dataset
ds = create_balanced_audio_dataset(
    sources={
        "speech": ["/data/speech"],
        "music": ["/data/music"],
    },
    sample_rate=44100,
    duration=3.0,
)

# Convert to IterDataset with multithreading
read_options = grain.ReadOptions(
    num_threads=4,  # 4 threads reading in parallel
    prefetch_buffer_size=16,  # Buffer 16 items per thread
)
iter_ds = ds.to_iter_dataset(read_options=read_options)

# Iterate with parallel loading
for audio in iter_ds:
    # Process audio
    pass

Threads help when the pipeline spends its time waiting (network storage, slow disks) and the Python work in between releases the GIL (audio decoding mostly happens in C extensions, so it usually does).

Multiprocessing with mp_prefetch

For CPU-bound work such as resampling, loudness, and augmentations, threads share one GIL and processes don’t, so use mp_prefetch:

import grain
from audiotree.sources import create_balanced_audio_dataset

# Create dataset
ds = create_balanced_audio_dataset(
    sources={
        "speech": ["/data/speech"],
        "music": ["/data/music"],
    },
    sample_rate=44100,
    duration=3.0,
)

# Convert to IterDataset and add multiprocessing
mp_options = grain.MultiprocessingOptions(
    num_workers=8,  # 8 worker processes
    per_worker_buffer_size=4,  # Each worker buffers 4 items
)
# Synchronous reads inside each worker: the processes are the parallelism.
read_options = grain.ReadOptions(num_threads=0, prefetch_buffer_size=0)
iter_ds = ds.to_iter_dataset(read_options=read_options).mp_prefetch(options=mp_options)

# Iterate with parallel loading
for audio in iter_ds:
    # Process audio
    pass

Combining Both

The two compose, with worker processes for the CPU work and a couple of threads inside each worker for I/O:

import grain
from audiotree.sources import create_balanced_audio_dataset

# Create dataset
ds = create_balanced_audio_dataset(
    sources={
        "speech": ["/data/speech"],
        "music": ["/data/music"],
    },
    sample_rate=44100,
    duration=3.0,
)

# Add multithreading for I/O
read_options = grain.ReadOptions(
    num_threads=2,
    prefetch_buffer_size=4,
)
iter_ds = ds.to_iter_dataset(read_options=read_options)

# Add multiprocessing for CPU-bound work
mp_options = grain.MultiprocessingOptions(
    num_workers=4,
    per_worker_buffer_size=2,
)
iter_ds = iter_ds.mp_prefetch(options=mp_options)

# Total threads: num_workers * num_threads = 4 * 2 = 8 threads
for audio in iter_ds:
    pass

Keep the multiplication in mind: this creates num_workers × num_threads threads in total (8 here), and it is easy to over-provision.

With Balanced Datasets

Weight-based balancing is preserved with multiprocessing:

from collections import Counter

import grain
from audiotree.sources import create_balanced_audio_dataset

# Create balanced dataset
ds = create_balanced_audio_dataset(
    sources={
        "speech": ["/data/speech"],
        "music": ["/data/music"],
    },
    weights={"speech": 0.7, "music": 0.3},
    sample_rate=44100,
    duration=3.0,
)

# Add multiprocessing
mp_options = grain.MultiprocessingOptions(num_workers=8)
read_options = grain.ReadOptions(num_threads=0, prefetch_buffer_size=0)
iter_ds = ds.to_iter_dataset(read_options=read_options).mp_prefetch(options=mp_options)

# Verify proportions are maintained
sources = [item.source[0] for item in iter_ds]
print(Counter(sources))
# {'speech': ~7000, 'music': ~3000}

Performance Tuning

Buffer Sizes

Buffers trade memory for throughput; bigger ones smooth over slow items but hold more decoded audio at once:

# Conservative (low memory)
mp_options = grain.MultiprocessingOptions(
    num_workers=4,
    per_worker_buffer_size=1,
)

# Aggressive (high throughput)
mp_options = grain.MultiprocessingOptions(
    num_workers=8,
    per_worker_buffer_size=8,
)

Worker Count

num_workers = os.cpu_count() is a reasonable ceiling to start from; reduce it if memory runs short. For an I/O-bound pipeline, threads are usually cheaper than more processes.

import os

# Use all CPU cores
num_workers = os.cpu_count()
mp_options = grain.MultiprocessingOptions(num_workers=num_workers)

Thread Count

0 (synchronous) is the baseline – see above. For I/O-bound workloads:

# Conservative
read_options = grain.ReadOptions(num_threads=2)

# Aggressive (for network storage)
read_options = grain.ReadOptions(num_threads=8)

Batching

Use AudioTree.batch with IterDataset.batch() to batch AudioTree objects. This concatenates along axis 0 (the batch dimension) rather than stacking, which would add an extra dimension.

import grain
from audiotree import AudioTree
from audiotree.sources import create_audio_dataset

ds = create_audio_dataset("/data/audio", duration=1.0)

# Batch with AudioTree.batch
iter_ds = ds.to_iter_dataset().batch(32, batch_fn=AudioTree.batch)

for batch in iter_ds:
    print(batch.waveform.shape)  # (32, channels, samples)

batch also handles dict structures containing AudioTrees:

# If your dataset yields {"src": AudioTree, "tgt": AudioTree}
iter_ds = ds.to_iter_dataset().batch(32, batch_fn=AudioTree.batch)

for batch in iter_ds:
    # batch is a dict with batched AudioTrees
    print(batch["src"].waveform.shape)  # (32, channels, samples)

See Working with Dict[str, AudioTree] Batches for more details on working with dict structures.

Typical Pipelines

Training Loop with Batching and Multiprocessing

import grain
import jax
from audiotree import AudioTree
from audiotree.sources import create_balanced_audio_dataset

# Create dataset
ds = create_balanced_audio_dataset(
    sources={"speech": ["/data/speech"], "music": ["/data/music"]},
    shuffle=True,
    num_epochs=None,  # Unbounded stream for training
    sample_rate=44100,
    duration=3.0,
)

# Add batching and multiprocessing
mp_options = grain.MultiprocessingOptions(
    num_workers=8,
    per_worker_buffer_size=4,
)
read_options = grain.ReadOptions(num_threads=0, prefetch_buffer_size=0)
iter_ds = (
    ds.to_iter_dataset(read_options=read_options)
    .batch(32, batch_fn=AudioTree.batch)
    .mp_prefetch(options=mp_options)
)

# Training loop
for step, batch in enumerate(iter_ds):
    if step >= max_steps:
        break

    # Convert to JAX array and train
    waveform = jnp.array(batch.waveform)  # (32, channels, samples)
    loss = train_step(waveform)

Validation with Deterministic Order

# Create validation dataset (no shuffle, exactly one epoch)
val_ds = create_balanced_audio_dataset(
    sources={"speech": ["/data/val_speech"], "music": ["/data/val_music"]},
    shuffle=False,  # Deterministic order
    num_epochs=1,
    excerpt_seed=42,  # Fixes which excerpt is drawn from each file
    sample_rate=44100,
    duration=3.0,
)

# Add batching and multiprocessing
mp_options = grain.MultiprocessingOptions(num_workers=4)
read_options = grain.ReadOptions(num_threads=0, prefetch_buffer_size=0)
val_iter_ds = (
    val_ds.to_iter_dataset(read_options=read_options)
    .batch(32, batch_fn=AudioTree.batch)
    .mp_prefetch(options=mp_options)
)

# Evaluate
for batch in val_iter_ds:
    metrics = evaluate(batch)

Worker Initialization

Perform setup in each worker process before loading data:

def worker_init_fn(worker_id, worker_count):
    """Called once per worker before processing data."""
    print(f"Worker {worker_id}/{worker_count} initialized")

    # Example: Set different random seed per worker
    import numpy as np

    np.random.seed(42 + worker_id)


mp_options = grain.MultiprocessingOptions(num_workers=4)
read_options = grain.ReadOptions(num_threads=0, prefetch_buffer_size=0)
iter_ds = ds.to_iter_dataset(read_options=read_options).mp_prefetch(
    options=mp_options,
    worker_init_fn=worker_init_fn,
)

Profiling

Enable profiling to identify bottlenecks:

mp_options = grain.MultiprocessingOptions(
    num_workers=8,
    per_worker_buffer_size=4,
    enable_profiling=True,  # Enable profiling
)

Logs will show timing information for each stage of the pipeline.

Memory Considerations

Decoded audio is big. Three seconds of 48 kHz stereo float32 is ~1.1 MB, a batch of 32 is ~35 MB, and 8 workers each buffering 4 such batches hold over a gigabyte between them before the model has allocated anything. When memory gets tight, shrink the buffers first, and use short durations while you’re still just testing the pipeline.

Troubleshooting

If workers hang or crash, check memory first: drop num_workers and per_worker_buffer_size and see whether the problem follows. If adding workers produces no speedup, the bottleneck is probably not the loader at all (the GPU, or network I/O); profile before tuning further. And if results change with the worker count, look for stateful operations like filter upstream of to_iter_dataset, since everything before that point should be deterministic.

More generally: start small (num_workers=2, per_worker_buffer_size=2), measure, and only then scale up. Parallelism added without measuring mostly adds memory pressure. If loading is still slow, check, in order: is mp_prefetch in place at all, are the buffers large enough to keep the GPU fed, are the files on slow storage, and is a saliency search enabled that you don’t actually need.

Example: Full Pipeline

Complete example with batching and all optimizations:

import grain
import jax
from audiotree import AudioTree
from audiotree.sources import create_balanced_audio_dataset
from audiotree.sources import ExcerptConfig

# Saliency for loud sections
excerpt = ExcerptConfig(
    strategy="loudest",
    lufs_cutoff=-40,
    num_tries=5,
)

# Create balanced dataset
ds = create_balanced_audio_dataset(
    sources={
        "speech": ["/fast_storage/speech"],
        "music": ["/fast_storage/music"],
        "effects": ["/fast_storage/effects"],
    },
    weights={"speech": 0.5, "music": 0.3, "effects": 0.2},
    shuffle=True,
    num_epochs=None,
    shuffle_seed=42,  # excerpt_seed defaults to this
    sample_rate=48000,
    duration=3.0,
    excerpt=excerpt,
)

# Synchronous reads: the mp_prefetch workers below are the parallelism.
# (Raise num_threads only for genuinely I/O-bound storage.)
read_options = grain.ReadOptions(num_threads=0, prefetch_buffer_size=0)
iter_ds = ds.to_iter_dataset(read_options=read_options)

# Batch with AudioTree.batch
iter_ds = iter_ds.batch(32, batch_fn=AudioTree.batch)

# Multiprocessing for CPU-bound work
mp_options = grain.MultiprocessingOptions(
    num_workers=8,
    per_worker_buffer_size=4,
    enable_profiling=False,  # Disable in production
)
iter_ds = iter_ds.mp_prefetch(options=mp_options)

# Training loop
for step, batch in enumerate(iter_ds):
    if step >= 100000:
        break

    # Your training code here
    waveform = jnp.array(batch.waveform)  # (32, channels, samples)
    loss = train_step(waveform)

    if step % 1000 == 0:
        print(f"Step {step}, Loss: {loss}")

See Also

• Chaining Transforms with Datasets - Chaining transforms with datasets

• Grain Documentation - Complete Grain pipeline guide

• Balanced Datasets - Creating balanced datasets

• create_audio_dataset() - Simple dataset creation

• create_balanced_audio_dataset() - Balanced dataset creation