Parallel And Larger-Than-Memory Workflows

Parallel And Larger-Than-Memory Workflows#

Raw binary outputs from simulations are often larger than the memory of a single machine. xarray-binfile builds on Dask, so the same code that reads a handful of files also scales to collections that do not fit in RAM.

This notebook shows how to:

  • open many files at once with open_mfdataset(..., parallel=True)

  • treat chunks as the unit of parallel, out-of-core work

  • choose a Dask scheduler and worker count for a compute step

  • evaluate several lazy results together with dask.compute

  • reduce data larger than memory without loading it all

  • stream a derived result back to disk one file at a time

Note

The dataset here is deliberately tiny so the documentation builds quickly. The important part is the pattern: because Dask keeps only a few chunks in memory at a time, the identical code scales to datasets much larger than RAM.

import pathlib
import tempfile

import dask
import numpy as np
import xarray as xr

import xarray_binfile  # noqa: F401  (registers the .binary_engine accessors)
from xarray_binfile.tutorial import DatasetGenerator, FileSpecsGetter

source_directory_holder = tempfile.TemporaryDirectory()
derived_directory_holder = tempfile.TemporaryDirectory()
source_directory = pathlib.Path(source_directory_holder.name)
derived_directory = pathlib.Path(derived_directory_holder.name)
source_directory

PosixPath('/tmp/tmpqv7lauuz')

Generate example data#

The source data uses float32 values on a small 3D grid, stored as one file per variable and timestep. This mirrors a typical simulation output directory with ux, uy, and uz velocity components.

base_coords = {
    "x": np.linspace(0.0, 3.0, num=32, dtype=np.float32),
    "y": np.linspace(-1.0, 1.0, num=24, dtype=np.float32),
    "z": np.linspace(0.0, 2.0, num=16, dtype=np.float32),
}

getter = FileSpecsGetter(base_coords=base_coords, dtype=np.float32)
dataset_generator = DatasetGenerator(getter.reader)
filenames = (
    getter.filename_template.format(name=name, digits=step)
    for name in ("ux", "uy", "uz")
    for step in range(12)
)
source_dataset = dataset_generator(map(pathlib.Path, filenames))
source_dataset.binary_engine.to_file(getter.writer, source_directory)
sorted(path.name for path in source_directory.glob("*.bin"))[:6]

['ux-0000.bin',
 'ux-0001.bin',
 'ux-0002.bin',
 'ux-0003.bin',
 'ux-0004.bin',
 'ux-0005.bin']

Open many files in parallel#

xr.open_mfdataset reads and combines the files by coordinate. Passing parallel=True opens the files concurrently with dask.delayed, and passing chunks keeps every variable lazy and Dask-backed instead of loading it eagerly.

lazy_dataset = xr.open_mfdataset(
    sorted(source_directory.glob("*.bin")),
    engine="binfile",
    read_specs_getter=getter.reader,
    chunks={"x": 8, "y": 6, "z": 4, "time": 3},
    parallel=True,
)
lazy_dataset

<xarray.Dataset> Size: 2MB
Dimensions:  (x: 32, y: 24, z: 16, time: 12)
Coordinates:
  * x        (x) float32 128B 0.0 0.09677 0.1935 0.2903 ... 2.71 2.806 2.903 3.0
  * y        (y) float32 96B -1.0 -0.913 -0.8261 -0.7391 ... 0.8261 0.913 1.0
  * z        (z) float32 64B 0.0 0.1333 0.2667 0.4 ... 1.6 1.733 1.867 2.0
  * time     (time) int64 96B 0 1 2 3 4 5 6 7 8 9 10 11
Data variables:
    ux       (x, y, z, time) float32 590kB dask.array<chunksize=(8, 6, 4, 1), meta=np.ndarray>
    uy       (x, y, z, time) float32 590kB dask.array<chunksize=(8, 6, 4, 1), meta=np.ndarray>
    uz       (x, y, z, time) float32 590kB dask.array<chunksize=(8, 6, 4, 1), meta=np.ndarray>

Chunks are the unit of parallel work#

Each chunk is an independent task that Dask can hand to a different worker, so more chunks means more opportunities for parallelism. Chunking is also what bounds memory use: only the chunks currently in flight are resident, no matter how large the full array is on disk.

ux = lazy_dataset["ux"]
blocks_per_dim = dict(zip(ux.dims, ux.data.numblocks))
largest_chunk_bytes = (
    int(np.prod([max(sizes) for sizes in ux.data.chunks])) * ux.dtype.itemsize
)

print("array type:", type(ux.data).__name__)
print("blocks per dim:", blocks_per_dim)
print("independent chunks (parallel tasks):", ux.data.npartitions)
print(f"whole ux array: {ux.nbytes / 1024:.1f} KiB")
print(f"largest single chunk: {largest_chunk_bytes / 1024:.1f} KiB")

array type: Array
blocks per dim: {'x': 4, 'y': 4, 'z': 4, 'time': 12}
independent chunks (parallel tasks): 768
whole ux array: 576.0 KiB
largest single chunk: 0.8 KiB

Choose a scheduler for a compute step#

Dask arrays default to the threaded scheduler, which suits the NumPy-heavy reads this backend performs. You can set the scheduler and worker count for a specific computation with dask.config.set. When running interactively, you can also wrap the call in Dask’s ProgressBar context manager to watch the chunks complete in parallel.

For CPU-bound pure-Python work use scheduler='processes', and for multi-machine clusters install dask.distributed and create a Client. The analysis code stays the same.

speed = np.sqrt(
    lazy_dataset["ux"] ** 2 + lazy_dataset["uy"] ** 2 + lazy_dataset["uz"] ** 2
).rename("speed")
lazy_mean_speed = speed.mean("time")

with dask.config.set(scheduler="threads", num_workers=4):
    mean_speed = lazy_mean_speed.compute()

mean_speed

<xarray.DataArray 'speed' (x: 32, y: 24, z: 16)> Size: 49kB
array([[[0.9239643 , 1.0230161 , 1.0118058 , ..., 1.0261247 ,
         0.8693645 , 1.0203921 ],
        [0.9672566 , 0.92100805, 0.9224105 , ..., 0.8687373 ,
         1.0682945 , 0.8341443 ],
        [1.0467513 , 0.94020873, 1.0428312 , ..., 1.0400995 ,
         1.045056  , 1.0226496 ],
        ...,
        [1.0616511 , 0.93975544, 0.7508979 , ..., 0.8954831 ,
         0.9754731 , 0.97374934],
        [1.0458447 , 1.0274231 , 0.9820626 , ..., 0.9773323 ,
         0.92604285, 0.9754898 ],
        [1.0712124 , 1.01089   , 1.0048076 , ..., 1.088391  ,
         0.73463565, 1.0104736 ]],

       [[0.8926453 , 1.0062106 , 1.0079592 , ..., 0.90749216,
         1.0042377 , 0.9367714 ],
        [0.9470346 , 0.8494458 , 0.84817713, ..., 0.7737827 ,
         0.8811388 , 0.87008953],
        [0.9777975 , 0.86743283, 0.96836954, ..., 0.9130836 ,
         1.0185399 , 0.9173892 ],
...
        [0.9277663 , 0.98586005, 0.95250124, ..., 1.1040207 ,
         0.9151638 , 1.0763098 ],
        [0.81757766, 1.0267109 , 0.9391477 , ..., 0.8497908 ,
         0.9602322 , 0.93783   ],
        [1.0143447 , 0.8509787 , 0.9242773 , ..., 1.0298214 ,
         0.9616863 , 0.90673184]],

       [[0.97884965, 0.9984443 , 1.0105523 , ..., 1.0137932 ,
         0.92400545, 0.8850007 ],
        [0.8248654 , 1.0559324 , 1.0184234 , ..., 0.9700162 ,
         0.8697305 , 1.1749941 ],
        [0.946641  , 1.015099  , 0.91480845, ..., 0.9958358 ,
         0.9408932 , 0.9630542 ],
        ...,
        [0.89720565, 1.0416396 , 1.0674402 , ..., 1.0540508 ,
         0.8764818 , 1.0257865 ],
        [0.87819785, 0.95982385, 0.9105374 , ..., 0.88527554,
         0.95447755, 1.0447162 ],
        [0.94379544, 0.9138942 , 0.86920184, ..., 0.8240814 ,
         0.892566  , 1.0514431 ]]], shape=(32, 24, 16), dtype=float32)
Coordinates:
  * x        (x) float32 128B 0.0 0.09677 0.1935 0.2903 ... 2.71 2.806 2.903 3.0
  * y        (y) float32 96B -1.0 -0.913 -0.8261 -0.7391 ... 0.8261 0.913 1.0
  * z        (z) float32 64B 0.0 0.1333 0.2667 0.4 ... 1.6 1.733 1.867 2.0

Evaluate several results together#

When you need more than one result, pass them all to dask.compute in a single call. Dask then builds one combined task graph, reads each shared chunk only once, and computes the outputs in parallel instead of walking the data multiple times.

lazy_mean = speed.mean("time")
lazy_max = speed.max("time")
lazy_std = speed.std("time")

mean_field, max_field, std_field = dask.compute(lazy_mean, lazy_max, lazy_std)
print("computed together:", mean_field.shape, max_field.shape, std_field.shape)
print(f"peak speed anywhere: {float(max_field.max()):.4f}")

computed together: (32, 24, 16) (32, 24, 16) (32, 24, 16)
peak speed anywhere: 1.7165

Reduce data larger than memory#

Reductions such as mean, std, or max consume the array chunk by chunk and keep only small partial results, so peak memory stays close to a single chunk rather than the whole dataset. This is why the same reduction works whether the collection is a few megabytes or hundreds of gigabytes, given enough time and disk.

overall_means = lazy_dataset.mean().compute()
print("domain- and time-averaged value per variable:")
overall_means

domain- and time-averaged value per variable:

<xarray.Dataset> Size: 12B
Dimensions:  ()
Data variables:
    ux       float32 4B 0.4997
    uy       float32 4B 0.5006
    uz       float32 4B 0.4999

Stream a derived result back to disk#

.binary_engine.to_file(...) materializes one sub-array (one file) at a time. With the default per-timestep writer, each timestep is computed, written, and released before the next one starts, so a derived dataset larger than memory can be persisted without ever holding it all at once.

The unit of streaming is the file, not the byte: each sub-array is fully loaded into memory (a Dask compute) and its file is written whole, in a single pass. There is no partial or append write mode, and each completed file is moved into place atomically with os.replace, which protects you from partially-updated, corrupted files even if the process is interrupted, but it also means every individual output file must fit in memory. Size the per-file slices in your write specs getter accordingly. If you need true streaming or partial writes, use one of the other xarray-supported formats such as NetCDF or Zarr.

speed.binary_engine.to_file(getter.writer, derived_directory)

written = sorted(path.name for path in derived_directory.glob("speed-*.bin"))
print(f"wrote {len(written)} files, one per timestep")
written[:6]

wrote 12 files, one per timestep

['speed-0000.bin',
 'speed-0001.bin',
 'speed-0002.bin',
 'speed-0003.bin',
 'speed-0004.bin',
 'speed-0005.bin']

Reading the derived files back and comparing against the in-memory result confirms the streamed write is correct.

roundtrip = xr.open_mfdataset(
    sorted(derived_directory.glob("speed-*.bin")),
    engine="binfile",
    read_specs_getter=getter.reader,
    chunks={"time": 4},
    parallel=True,
)
xr.testing.assert_allclose(roundtrip["speed"].compute(), speed.compute())
print("round-trip matches the in-memory speed field")
roundtrip

round-trip matches the in-memory speed field

<xarray.Dataset> Size: 590kB
Dimensions:  (x: 32, y: 24, z: 16, time: 12)
Coordinates:
  * x        (x) float32 128B 0.0 0.09677 0.1935 0.2903 ... 2.71 2.806 2.903 3.0
  * y        (y) float32 96B -1.0 -0.913 -0.8261 -0.7391 ... 0.8261 0.913 1.0
  * z        (z) float32 64B 0.0 0.1333 0.2667 0.4 ... 1.6 1.733 1.867 2.0
  * time     (time) int64 96B 0 1 2 3 4 5 6 7 8 9 10 11
Data variables:
    speed    (x, y, z, time) float32 590kB dask.array<chunksize=(32, 24, 16, 1), meta=np.ndarray>

Where to go next#

The scaling behavior here comes from Dask and xarray directly, so their guides apply unchanged:

For portable archival of results, prefer Dataset.to_netcdf(...) or Dataset.to_zarr(...); reach for raw binary output only when an external tool requires that exact layout.