Scaling Raster Processing with Dask

A single Sentinel-2 tile is a few hundred megabytes; a multi-year stack of them is hundreds of gigabytes. The moment your analysis window exceeds the RAM of one process, naive NumPy loads fail with a MemoryError, and looping scene-by-scene throws away the parallelism modern machines offer. Dask solves both problems at once: it splits arrays into chunks, builds a lazy task graph over them, and executes that graph out-of-core across a pool of workers so peak memory stays bounded no matter how large the cube grows. For where this fits in the broader deployment picture — orchestration, distributed clusters, and cost tuning — see the parent section, Cloud Execution & Orchestration.

This guide is for remote sensing analysts and data engineers who already compute indices on single scenes and now need to run the same math over a (time, band, y, x) cube that will not fit in memory. It covers cluster setup, lazy stack reads, mapping band math across chunks, the persist versus compute decision, dashboard monitoring, and streaming results to disk.


Execution model and architecture

Dask separates graph construction from execution. Opening a raster with chunks= and applying arithmetic only records operations; nothing runs until you call .compute(), .persist(), or a write. The diagram below shows how a lazy xarray graph is scheduled onto workers by the distributed Client.

Dask out-of-core raster execution architecture A chunked xarray cube feeds a lazy task graph. The distributed Client and scheduler dispatch chunk-level tasks to several worker processes, each holding one chunk in memory at a time and spilling to local disk under pressure, then streaming completed blocks to output storage. Chunked cube (time, band, y, x) Lazy task graph open · scale · (NIR−Red) /(NIR+Red) · median built, not yet run Client + Scheduler dispatch · dashboard Worker 1 1 chunk in RAM Worker 2 spill → disk Worker N threads × procs Output Zarr / COG Nothing executes until compute/persist/write triggers the graph; workers hold one chunk at a time LocalCluster when… fits one machine · cores < 64 data spills to local SSD interactive / single-node batch Distributed cluster when… RAM exceeds one node workers co-located with S3 Coiled / AWS Batch / Kubernetes

The right-hand decision band previews a recurring theme: a LocalCluster is the correct default until the data or core count outgrows one machine, at which point the same code runs against a distributed cluster launched on Coiled & AWS Batch. Only the Client connection string changes.


Prerequisites

Install the stack, pinning majors so a GDAL or dask update does not silently change chunk semantics or scheduler behaviour:

pip install "dask[distributed,array]>=2024.1" "xarray>=2024.1" \
            "rioxarray>=0.15" "rasterio>=1.3" "stackstac>=0.5" \
            "zarr>=2.16" "bokeh>=3.1"
Library Min version Why required
dask[distributed,array] 2024.1 Chunked arrays, task graph, LocalCluster, Client
xarray 2024.1 Labeled (time, band, y, x) cubes with lazy Dask backing
rioxarray 0.15 CRS-aware raster reads returning Dask-backed DataArrays
rasterio 1.3 GDAL I/O, windowed and /vsicurl/ reads under xarray
stackstac 0.5 Turn a STAC item list into one lazy chunked cube
zarr 2.16 Chunk-aligned out-of-core writes without materialising the cube
bokeh 3.1 Renders the live scheduler dashboard

Conceptual prerequisites: You should already be comfortable computing indices on a single array — see Band Math Operations with xarray — and know how to assemble a scene list from a catalog via Querying STAC Catalogs Programmatically. Understanding Cloud-Optimized GeoTIFF internal tiling helps you align Dask chunks to the file’s own blocks.


Step-by-step workflow

Step 1 — Install the Dask and xarray stack

Confirm the distributed scheduler imports cleanly and that Bokeh is present so the dashboard will render. A missing bokeh is the most common reason the diagnostics URL 404s:

import dask
import xarray as xr
from dask.distributed import LocalCluster, Client

print(f"dask {dask.__version__}, xarray {xr.__version__}")

try:
    import bokeh  # noqa: F401
    print("bokeh present — dashboard will render")
except ImportError:
    raise EnvironmentError(
        "Install bokeh to enable the scheduler dashboard: pip install bokeh"
    )

Step 2 — Start a LocalCluster and connect a Client

A LocalCluster spins up worker processes on the current machine; the Client is your handle to submit graphs and open the dashboard. Size workers deliberately — for chunked raster math, a few processes each with several threads balances GIL-releasing NumPy calls against memory isolation:

from dask.distributed import LocalCluster, Client

cluster = LocalCluster(
    n_workers=4,               # separate processes → isolated memory, no GIL contention
    threads_per_worker=2,      # threads share a heap; good for GIL-releasing NumPy/GDAL
    memory_limit="4GB",        # per-worker cap; spill starts ~70%, kill near 95%
    dashboard_address=":8787",
)
client = Client(cluster)
print(client.dashboard_link)  # open this in a browser to watch tasks live

Keep the total (n_workers × memory_limit) below physical RAM with headroom for the OS page cache — GDAL leans on it heavily for /vsicurl/ reads. To target a distributed backend later, replace these two lines with Client("tcp://scheduler:8786") and the rest of the notebook is unchanged.

Step 3 — Open a raster stack lazily with chunks

Reading with chunks= returns a Dask-backed DataArray: coordinates and dtype are known immediately, but no pixels move yet. Chunk the spatial dimensions and keep bands whole so per-pixel index math needs no cross-chunk communication:

import rioxarray

# One scene, chunked spatially; bands kept whole (band=-1) for per-pixel math
da = rioxarray.open_rasterio(
    "s2_scene_B0408.tif",
    chunks={"band": -1, "y": 2048, "x": 2048},
    lock=False,                # allow concurrent GDAL reads across threads
)
print(da)                       # dask.array in the repr — still lazy
print(da.data.chunksize)        # e.g. (2, 2048, 2048)

For a multi-scene cube, stackstac.stack(items, chunksize=2048) returns a single (time, band, y, x) DataArray spanning every STAC item. Chunk sizing is the single biggest performance lever here; the dedicated guide, Tuning Dask Chunk Sizes for Raster Cubes, works through the 256–512 MB target and the time=1, band=-1 reasoning in full.

Step 4 — Map index math across chunks lazily

Ordinary xarray arithmetic on a Dask-backed array extends the graph rather than computing. Because each output pixel depends only on the co-located input pixels, the scheduler runs one chunk per task with no shuffle:

import numpy as np

nir = da.sel(band=8).astype("float32")
red = da.sel(band=4).astype("float32")

denom = nir + red
ndvi = xr.where(denom != 0, (nir - red) / denom, np.nan)
ndvi = ndvi.rename("ndvi")

print(ndvi.data.npartitions, "tasks queued")  # still nothing computed

The xr.where guard keeps division-by-zero out of the result and matches the numerical-stability pattern used throughout Spectral Index Calculation Pipelines. The end-to-end multi-scene version — stack, lazy NDVI, median composite — is covered in Computing NDVI with xarray on a Dask Cluster.

Step 5 — Choose persist versus compute

.compute() runs the graph and pulls the entire result into the client process as NumPy — only safe when the result is small (a scalar, a time series, one composite). .persist() runs the graph but leaves the result distributed across worker memory as chunks, so downstream steps reuse it without re-reading source pixels:

# Pin the expensive stack in cluster RAM once, reused by every branch below
ndvi = ndvi.persist()
from dask.distributed import wait
wait(ndvi)                      # block until all chunks are materialised on workers

# Cheap reductions now read from RAM, not from disk/S3
mean_curve = ndvi.mean(dim=("y", "x")).compute()   # small → safe to pull local
p90 = ndvi.quantile(0.9, dim="time")               # still lazy; another branch

The rule of thumb: persist what you will reuse, compute only what is small enough to hold locally. Persisting an array larger than aggregate worker RAM triggers spilling and defeats the purpose — size it against the same budget you use for chunks.

Step 6 — Monitor the task graph on the dashboard

Open client.dashboard_link (default http://localhost:8787) while a computation runs. The Task Stream shows per-worker execution; the Progress bars show completion by task type; the Workers tab shows live memory per process. Use it to diagnose in real time:

# Trigger a real computation, then watch :8787/status
result = ndvi.mean(dim=("y", "x")).compute()

# Programmatic snapshot without the browser
info = client.scheduler_info()
for wid, w in info["workers"].items():
    print(wid, "mem", round(w["metrics"]["memory"] / 1e9, 2), "GB")

Red blocks in the Task Stream mean spilling-to-disk; a mostly-white stream means workers are starved (chunks too large or too few). Solid colour with little white is a healthy, well-chunked graph.

Step 7 — Write results without materialising the whole cube

Never call .compute() on a cube-sized result and hand it to a writer — that reassembles everything in the client. Instead stream chunk-aligned to Zarr, or write per-scene COGs. Zarr writes each Dask chunk independently and in parallel:

# Align the on-disk Zarr chunking to the Dask chunking to avoid rewrite overhead
ndvi.to_dataset(name="ndvi").chunk({"y": 2048, "x": 2048}).to_zarr(
    "ndvi_cube.zarr",
    mode="w",
    consolidated=True,          # single metadata read on reopen
)

For a georeferenced GeoTIFF composite, reduce the time dimension first so the result is a single 2-D array, then rio.to_raster(..., tiled=True, compress="zstd"). Building a browser-servable Cloud-Optimized GeoTIFF adds overviews on top of that write.


Parameter reference

Parameter Type Default Usage note
n_workers int cores Number of worker processes. Fewer, fatter workers reduce inter-process transfer for spatial math
threads_per_worker int auto Threads per process; 2–4 works well since GDAL/NumPy release the GIL
memory_limit str auto Per-worker cap. Spill ~0.7, pause ~0.8, terminate ~0.95 of this value
chunks (open) dict file blocks Chunk shape per dim. {"time":1,"band":-1,"y":2048,"x":2048} is the raster default
chunksize (stackstac) int 1024 Square spatial chunk edge for a STAC-derived cube
lock bool True Set False to allow concurrent GDAL reads across threads on thread-safe drivers
.persist() method Materialise + keep distributed in worker RAM; use for reused intermediates
.compute() method Materialise + pull to client as NumPy; only for small results
consolidated bool False Zarr metadata consolidation; True cuts reopen latency on object stores

Verification & testing

Confirm the array is genuinely lazy before you compute, and that a small reduction returns finite values:

import dask.array as darr
import numpy as np

# 1. The array must be Dask-backed, not eager NumPy
assert isinstance(ndvi.data, darr.Array), "Array is not chunked — reads will be eager"

# 2. Chunk bytes should sit in the target band (see the chunk-tuning guide)
nbytes_per_chunk = np.prod(ndvi.data.chunksize) * ndvi.dtype.itemsize
assert nbytes_per_chunk < 600e6, f"Chunk too large: {nbytes_per_chunk/1e6:.0f} MB"

# 3. A cheap reduction must run and stay in range
sample = ndvi.isel(x=slice(0, 512), y=slice(0, 512)).mean().compute()
assert np.isfinite(sample), "Reduction returned NaN/Inf — check the divide guard"
print(f"OK — chunk {nbytes_per_chunk/1e6:.0f} MB, sample mean {float(sample):.3f}")

Run this against a spatial slice, not the full cube, so verification itself stays cheap. In CI, point it at a small fixture scene and assert the chunk-size and finiteness invariants.


Troubleshooting

distributed.nanny - WARNING - Worker exceeded 95% memory budget. Restarting — chunks are too large for memory_limit, or you persisted more than fits. Shrink chunks toward a 256 MB target, cut the number of live persisted arrays, or raise memory_limit. The chunk-tuning guide has the sizing math.

Computation is slow and the Task Stream is mostly white — workers are starved because chunks are too big (few coarse tasks) or the graph has a serial bottleneck. Reduce chunk edge length to expose more parallel tasks, and avoid .compute() inside a loop.

KilledWorker on a specific task — one chunk blows the memory budget mid-task, usually from an accidental .load() or a groupby that shuffles. Inspect the failing key in the dashboard, and prefer xr.where/reductions that stay chunk-local over operations that rechunk.

Dashboard link 404s or refuses connectionbokeh is not installed, or the port is taken. Install bokeh and set an explicit dashboard_address=":8787"; on a remote host, SSH-forward the port.

Every .compute() re-reads from S3 — you never persisted the shared intermediate. Call .persist() once on the reused array and wait() for it before branching multiple outputs off it.

ValueError: conflicting sizes for dimension when stacking scenes — the input scenes have different grids. Reproject or resample to a common grid first (see Advanced Resampling & Upscaling Techniques) before building the cube.


FAQ

When should I move from a LocalCluster to a distributed cluster?

Stay on a LocalCluster while the workload fits one machine’s cores and its intermediates spill comfortably to local SSD. Move to a multi-machine cluster when a single node’s aggregate RAM cannot hold even a few persisted intermediates, when you need more cores than one box provides, or when the data lives in object storage and you want workers in the same cloud region to cut read latency and egress.

Does calling .compute() twice re-read the data from disk?

Yes. Each .compute() re-executes the whole graph from its source reads unless you pinned an intermediate with .persist(). If several outputs branch off one expensive lazy result, persist that result once, wait() for it, then compute each branch so the shared upstream work runs a single time.

Why does my Dask computation spill to disk or get killed?

The chunks are too large for the per-worker memory budget, or too many intermediates are held live at once. Dask targets a few hundred megabytes per chunk; past ~70% of memory_limit a worker spills, and near 95% the nanny kills it. Shrink the chunks, reduce persisted arrays, or raise memory_limit.


Deep-Dive Articles