Tuning Dask Chunk Sizes for Raster Cubes

Target 256–512 MB of uncompressed data per chunk, chunk the spatial dimensions while keeping band=-1 and time=1, and verify the byte size with a one-line nbytes calculation before you compute anything. For a float32 (time, band, y, x) cube that lands on a roughly 2048×2048 spatial tile:

import numpy as np, rioxarray

da = rioxarray.open_rasterio(
    "s2_cube.tif",
    chunks={"band": -1, "y": 2048, "x": 2048},  # spatial-first; bands whole
)
chunk_mb = np.prod(da.data.chunksize) * da.dtype.itemsize / 1e6
print(f"{chunk_mb:.0f} MB per chunk")            # keep this in the 256–512 band

This is the single highest-leverage setting when scaling raster processing with Dask: get it right and the graph runs smoothly; get it wrong and you either drown in scheduler overhead or blow the worker memory budget.

Why This Arises When Scaling Raster Work

Dask has to pick a granularity. Every chunk becomes one task in the graph, and the scheduler carries a fixed per-task cost — roughly a millisecond of bookkeeping plus the memory to track the key. Make chunks tiny and a continental cube explodes into millions of tasks whose overhead dwarfs the actual arithmetic. Make chunks huge and each task needs gigabytes of RAM, so workers spill to disk or get killed by the nanny. The sweet spot balances these two failure modes, and for raster cubes it sits in the hundreds-of-megabytes range.

The complication specific to a (time, band, y, x) cube is that the four dimensions are not interchangeable. Bands are read together and consumed together by index math; scenes come from separate files at separate timestamps; the spatial dimensions are where the data volume lives. So chunk sizing is not one number but a shape decision, and the right shape follows the access pattern. This page focuses narrowly on that shape. For cluster setup, persist versus compute, and writing results, see the parent guide, Scaling Raster Processing with Dask, inside the broader Cloud Execution & Orchestration section.

The chunk shape, visualised

The diagram shows how a (time, band, y, x) cube is sliced: whole in band, one step in time, tiled in y and x. Each shaded block is one chunk — one task the scheduler dispatches.

Chunking a (time, band, y, x) raster cube The spatial y and x plane is divided into a grid of tiles; each tile spans all bands and one time step, forming a chunk of roughly 256 to 512 megabytes. Time steps stack behind as separate chunks read from separate scene files. 1 chunk 2048×2048 all bands x (tiled) y (tiled) time (chunk = 1) one scene per chunk Chunk shape rules band = -1 → whole, math needs all bands time = 1 → one scene, read per file y, x = 2048 → tile the volume dimensions target = 256–512 MB uncompressed / chunk bytes = ∏(chunk dims) × itemsize 6 × 2048 × 2048 × 4 B ≈ 100 MB too small → millions of tasks, scheduler overhead too large → spill to disk, KilledWorker

Environment & Setup

Only three packages are needed to reason about and set chunk sizes:

Package Minimum version Why required
dask 2024.1 Chunked array engine and the .rechunk operation
xarray 2024.1 Labeled cube whose .chunk() accepts per-dimension sizes
rioxarray 0.15 Opens rasters directly into a Dask-backed DataArray with chunks= 
pip install "dask[array]>=2024.1" "xarray>=2024.1" "rioxarray>=0.15"

Complete Working Example

The helper below picks a spatial chunk edge that lands inside the 256–512 MB target for a given cube shape and dtype, then reports the resulting task count so you can sanity-check granularity before launching a job:

import numpy as np
import xarray as xr


def choose_spatial_chunk(
    da: xr.DataArray,
    target_mb: float = 384.0,     # midpoint of the 256–512 MB band
    edge_step: int = 256,         # round to multiples of the COG internal tile
) -> dict[str, int]:
    """
    Return a chunks dict for a (time, band, y, x) DataArray.

    Keeps band whole (-1) and time at 1, then sizes the square spatial
    edge so one chunk's uncompressed bytes sit near target_mb.
    """
    itemsize = da.dtype.itemsize                  # float32 → 4 bytes
    n_band = da.sizes.get("band", 1)              # all bands live in one chunk

    # bytes = n_band * edge * edge * itemsize  →  solve edge for target_mb
    target_bytes = target_mb * 1e6
    raw_edge = (target_bytes / (n_band * itemsize)) ** 0.5

    # Round DOWN to a multiple of edge_step so chunks align to COG tiles
    edge = max(edge_step, int(raw_edge // edge_step) * edge_step)
    # Never exceed the actual array extent
    edge = min(edge, da.sizes["y"], da.sizes["x"])

    return {"time": 1, "band": -1, "y": edge, "x": edge}


def report_chunks(da: xr.DataArray) -> None:
    """Print per-chunk size and total task count for a chunked DataArray."""
    cs = da.data.chunksize
    mb = np.prod(cs) * da.dtype.itemsize / 1e6
    print(f"chunk shape {cs} → {mb:.0f} MB/chunk, {da.data.npartitions} tasks")


if __name__ == "__main__":
    # A synthetic 12-scene, 6-band, 10980×10980 Sentinel-2-sized float32 cube
    shape = (12, 6, 10980, 10980)
    lazy = xr.DataArray(
        __import__("dask.array", fromlist=["zeros"]).zeros(shape, dtype="float32"),
        dims=("time", "band", "y", "x"),
    )
    chunks = choose_spatial_chunk(lazy)
    print("chosen chunks:", chunks)               # e.g. {'time':1,'band':-1,'y':3840,'x':3840}
    report_chunks(lazy.chunk(chunks))

The edge_step rounding matters: aligning chunk edges to the raster’s internal 256- or 512-pixel Cloud-Optimized GeoTIFF tiles means each Dask read maps to whole file blocks, avoiding partial-block reads that re-fetch the same bytes.

Variant Patterns

1. Rechunking before a temporal reduction

A median or percentile over time with time=1 forces Dask to gather every scene for each spatial tile — many small tasks feeding one reduction. Widen the spatial chunk first so there are fewer, fatter reduction tasks:

# Fewer spatial tiles → fewer time-gather tasks in the median
composite = (
    cube.chunk({"time": -1, "y": 1024, "x": 1024})   # collapse time, shrink space
        .median(dim="time")
)

Rechunk deliberately and once: rechunking moves data between workers over the network and is one of the most expensive things you can ask a Dask graph to do.

2. Diagnosing chunks that are too small

If report_chunks shows hundreds of thousands of tasks and the dashboard Task Stream is a blizzard of sub-second blocks, the chunks are too fine. Symptoms are high scheduler CPU, slow graph construction (not execution), and memory growth on the scheduler itself:

# Merge neighbouring tiles to cut task count 4× without changing the result
coarser = fine_cube.chunk({"y": 4096, "x": 4096})
report_chunks(coarser)   # verify it stays under ~512 MB after merging

3. Diagnosing chunks that are too large

The opposite symptom — a worker jumping to 90%+ memory on a single task, red spill blocks, or KilledWorker — means a chunk exceeds the per-worker budget. Halve the edge and recheck the byte size:

mb = np.prod(cube.data.chunksize) * cube.dtype.itemsize / 1e6
if mb > 512:
    cube = cube.chunk({"y": cube.chunksizes["y"][0] // 2,
                       "x": cube.chunksizes["x"][0] // 2})

Chunk sizing and worker memory_limit are two sides of one budget; a chunk should be well under the RAM one worker thread can hold, because several intermediates coexist during a computation. Once sizing is right, the end-to-end pattern in Computing NDVI with xarray on a Dask Cluster runs without spilling.

Common Errors

distributed.worker - WARNING - Memory use is high but worker has no data to store to disk One chunk is larger than the memory a single task can hold. Reduce the spatial edge length so per-chunk bytes drop below ~256 MB, and confirm with the nbytes formula.

Graph takes seconds to build before any compute starts Too many chunks — the task count is in the hundreds of thousands. Increase the spatial edge (or collapse an over-fine time chunking) so npartitions drops into the thousands.

PerformanceWarning: Slicing is producing a large chunk An operation is unifying chunks along a dimension you split, often an .sel across time. Rechunk that dimension to -1 deliberately before the operation rather than letting Dask do it implicitly mid-graph.