Tuning Dask Chunk Size for Spatial Tiling

Chunk size is the single knob that decides whether a tiled Dask generation job finishes in minutes or thrashes for hours, and the right value is a narrow band bounded below by scheduler overhead and above by per-worker memory.

Part of Async Execution for Large Grids: that workflow covers distributing generation across a worker pool; this page addresses the concrete tuning decision — how large each spatial chunk should be — and the two ways an unconsidered choice degrades a job: a task graph so large the scheduler becomes the bottleneck, or chunks so heavy that workers spill and stall.

Root Cause: Chunk Size Sits Between Two Opposing Costs

A Dask array over a large grid is a lazy graph of tasks, one per chunk. Every operation you apply multiplies that graph. The wall-clock cost of a tiled job is dominated by two terms that move in opposite directions as chunk size changes.

Scheduler and per-task overhead (grows as chunks shrink). Each task carries fixed cost: graph construction, serialization of the task spec, scheduling, and result bookkeeping. Dask’s own guidance puts useful task duration in the range of roughly 100 ms to 1 s; below that, the scheduler’s per-task overhead (order of hundreds of microseconds to a millisecond each) starts to dominate. A 100k×100k grid cut into 256×256 chunks is over 152,000 tasks per operation — the scheduler serializes and tracks a graph so large it can exhaust the client’s memory before a single worker starts real work.

Per-worker memory pressure (grows as chunks grow). A worker must hold, simultaneously, every chunk an in-flight task touches plus its output. A neighbourhood or halo operation holds the chunk and its overlap. Dask targets keeping each chunk small enough that many fit in a worker’s memory at once — a common rule of thumb is chunks in the low hundreds of MB, well under the per-worker memory limit divided by the thread count. Overshoot and the worker spills to disk, then pauses, then in the worst case is killed and its tasks are recomputed elsewhere — the pathology that turns a slow job into a stuck one.

The target is the band where a single chunk’s task runs long enough to amortize overhead but small enough that the worker’s thread pool can hold several concurrently. For a float64 raster, a chunk of c×cc \times c cells occupies 8c28 c^2 bytes; a 2048×2048 float64 chunk is 32 MiB, a comfortable size that yields tasks in the hundreds-of-milliseconds range for typical per-cell generation work. This is also why the chunking cannot silently change generated values: seeds must be derived from spatial tile identity, not from partition index, so re-tuning chunk size for performance never perturbs the output.

U-shaped total job time versus Dask chunk size with an optimal band between task overhead and memory spill A chart with chunk size on the horizontal axis increasing to the right, and total job time on the vertical axis increasing upward. A descending curve labelled scheduler and per-task overhead dominates at small chunk sizes on the left. An ascending curve labelled per-worker memory spill and recompute dominates at large chunk sizes on the right. Their sum is a U-shaped total-time curve with a broad minimum. A shaded vertical band marks the optimal region, roughly a chunk of 2048 by 2048 float64 cells at about 32 mebibytes, aligned to the raster block size. The left edge is annotated too many tiny tasks, graph exceeds client memory; the right edge is annotated worker spill, pause, and kill. Total job time Chunk size → optimal band ~2048² f64 ≈ 32 MiB scheduler / per-task overhead worker spill / recompute total = overhead + memory too many tiny tasks; graph > client memory spill → pause → kill
Total time is the sum of a falling overhead term and a rising memory term; the minimum is a band, not a point, so aim for it and align to the raster block.

Prerequisite Check: Inspect the Task Graph and Chunk Bytes

Before tuning, measure what you have. Pin the toolchain and read the chunk layout and task count directly rather than guessing.

dask==2024.*
numpy==1.26.*
rasterio==1.3.*     # to read the on-disk block/tile size

python
import numpy as np
import dask.array as da

# A synthetic 100k x 100k float64 grid cut too finely.
arr = da.random.random((100_000, 100_000), chunks=(256, 256))

print("chunk bytes:", arr.chunksize[0] * arr.chunksize[1] * arr.dtype.itemsize)  # 524288 -> 0.5 MiB
print("num chunks :", arr.npartitions)                                           # ~152881
print("est tasks  :", len((arr + 1).__dask_graph__()))                           # scales with npartitions

A chunk under a few MiB combined with a partition count in the hundreds of thousands is the signature of over-fine tiling: the graph is enormous and each task does trivial work. The fix is to re-chunk toward the target band and align to the storage block.

Fix: Size Chunks to a Target, Align to Blocks, Bound the Halo

Set an explicit byte target, snap chunk dimensions to a multiple of the on-disk block size so a chunk read never touches partial blocks, and — for neighbourhood operations — bound the halo so overlap does not silently multiply memory. This tiling feeds directly into memory-optimized large-grid generation, which assumes each chunk fits comfortably in a worker.

python
import math
import numpy as np
import dask.array as da

def target_chunk_edge(dtype=np.float64, target_mib=32, block=256) -> int:
    """Square chunk edge (in cells) near a byte target, snapped to the block size.
    Aligning to the raster block avoids read amplification at chunk boundaries."""
    itemsize = np.dtype(dtype).itemsize
    raw_edge = math.sqrt(target_mib * 2**20 / itemsize)   # cells per side for target bytes
    edge = max(block, round(raw_edge / block) * block)    # snap to a multiple of block
    return int(edge)

edge = target_chunk_edge(target_mib=32, block=256)        # -> 2048 (2048^2 * 8 B = 32 MiB)
arr = da.random.random((100_000, 100_000), chunks=(edge, edge))
print(arr.npartitions, "chunks;", arr.chunksize[0]*arr.chunksize[1]*8 // 2**20, "MiB each")

For a windowed/neighbourhood generator, use map_overlap and keep the halo depth as small as the stencil truly requires. Every extra cell of depth is loaded and recomputed on all four sides of every interior chunk:

python
import dask.array as da

def generate_with_halo(arr: da.Array, radius: int) -> da.Array:
    """Neighbourhood generation with a minimal halo.
    Overlap memory per chunk scales with (edge + 2*radius)^2, so keep radius tight."""
    return arr.map_overlap(
        lambda block: neighbourhood_kernel(block, radius),
        depth=radius,            # ONLY the stencil radius, never 'to be safe' padding
        boundary="reflect",      # define the seam explicitly; avoid NaN bleed at the extent
        trim=True,               # drop the halo from the result so tiles are seamless
    )

The halo-recompute trap is the most expensive mistake here: a depth set generously “to be safe” is not free — with a 2048 chunk and depth=64, each interior chunk actually materializes 2176×2176 cells, a 13% memory and compute surcharge that compounds across every operation in the graph. Set depth to the exact stencil radius and no more.

Verification Step: Assert Task Count and Chunk Bytes Stay in Band

Gate the tuning so a future edit cannot silently re-introduce over-fine chunks or oversized ones. Bound both the per-chunk byte size and the total task count.

python
import numpy as np
import dask.array as da

def test_chunking_is_in_band():
    edge = target_chunk_edge(target_mib=32, block=256)
    arr = da.random.random((100_000, 100_000), chunks=(edge, edge))

    # 1. Per-chunk bytes within the target band (8 MiB .. 128 MiB).
    chunk_bytes = arr.chunksize[0] * arr.chunksize[1] * arr.dtype.itemsize
    assert 8 * 2**20 <= chunk_bytes <= 128 * 2**20, f"{chunk_bytes} B out of band"

    # 2. Chunk edge aligned to the on-disk block size (no read amplification).
    assert arr.chunksize[0] % 256 == 0, "chunk not block-aligned"

    # 3. Task graph small enough for the client to hold.
    assert arr.npartitions < 10_000, f"{arr.npartitions} chunks -> graph too large"

    # 4. Halo surcharge bounded: overlap must add < 15% area for a radius-64 stencil.
    edge_with_halo = edge + 2 * 64
    assert (edge_with_halo**2) / (edge**2) < 1.15, "halo depth too large for chunk size"

The task-count ceiling is the assertion that most often fires in practice: it is easy to shrink chunks for “more parallelism” and accidentally cross back over the scheduler-overhead cliff. Tie the ceiling to your cluster’s client memory rather than a magic number when you productionize.

Edge Cases & Gotchas

Chunks that straddle the antimeridian. When a global grid in EPSG:4326 is tiled by longitude, the chunk containing ±180° wraps, and a map_overlap with boundary="reflect" reflects across the seam instead of wrapping, producing a discontinuity at the dateline. For global extents, tile in a projection whose x-axis does not wrap (an equal-area CRS such as EPSG:6933) or set boundary="periodic" on the longitude axis so the halo pulls from the opposite edge.

Misaligned on-disk blocks. If the source Cloud-Optimized GeoTIFF is written in 512×512 internal tiles but you chunk in 2048×2048 that is not a multiple of 512 at the array edge, the final partial chunk forces partial-block reads and read amplification. Query the block size with rasterio (ds.block_shapes) and snap chunk edges to a multiple of it; this is the same alignment contract that streaming large grids to Cloud-Optimized GeoTIFF enforces on write.

Uneven final chunks skewing the memory estimate. A grid whose dimensions are not a multiple of the chunk edge leaves thin edge chunks. Dask sizes memory by the largest chunk, so a single oversized remainder chunk can trigger spills even when the average chunk is comfortable. Use arr.rechunk with an explicit block scheme, or pad the grid to a multiple of the edge and crop after generation, so every chunk is uniform.

Frequently Asked Questions

What chunk size should I start with for float64 raster generation?
Start near a 32 MiB chunk — a 2048 by 2048 float64 tile — and snap the edge to a multiple of the source's on-disk block size. That puts per-task work in the hundreds-of-milliseconds range, keeps the task graph small enough for the client, and leaves headroom for several chunks per worker thread. Adjust from there using the dashboard's task-stream and memory panels.
Why is my job spilling to disk even though each chunk looks small?
Dask sizes worker memory by the largest chunk, and a neighbourhood operation holds the chunk plus its halo simultaneously. An oversized remainder chunk on a non-uniform grid, or a generous map_overlap depth, can push a single task well past the average chunk size. Rechunk to uniform blocks and set the halo depth to exactly the stencil radius.
Does changing chunk size change the generated values?
It must not. If it does, seeds are being derived from partition index or arrival order instead of spatial tile identity. Derive each tile's random stream from its spatial coordinates so re-chunking for performance never perturbs output, and gate the pipeline on a content hash to prove it.
How small is too small for a chunk?
When per-task work drops below roughly 100 milliseconds, the scheduler's per-task overhead starts to dominate, and when the partition count reaches the hundreds of thousands the task graph can exhaust the client process before workers begin. A chunk under a few MiB on a large grid is usually a sign you have crossed that line; coarsen toward the target band.