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.
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 connection — bokeh 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.
Related
- Cloud Execution & Orchestration — parent section covering orchestration, distributed clusters, and pipeline cost tuning
- Tuning Dask Chunk Sizes for Raster Cubes — pick chunk shapes that balance task overhead against worker memory
- Computing NDVI with xarray on a Dask Cluster — an end-to-end multi-scene NDVI composite on a distributed client
- Distributed Processing on Coiled & AWS Batch — run the same graphs on a managed multi-node cluster
- Band Math Operations with xarray — the single-array index math that Dask parallelises
- Spectral Index Calculation Pipelines — the production index workflow this scales out