Parallel Tile Processing with Dask

TL;DR: Start a memory-bounded LocalCluster, open the orthomosaic as a lazily-chunked dask.array (or xarray DataArray) aligned to its internal tiling, apply NDVI blockwise with map_blocks, and write the result only at the final compute — then assert the output shape and CRS match the input.

Why This Task Arises

A stitched multispectral orthomosaic of a large field is routinely 8–30 GB across five bands. Computing NDVI by reading the whole red and NIR bands into NumPy arrays needs the full raster in RAM twice over, and on a worker with 16 GB that is an immediate MemoryError. The naive fix — a hand-written windowed loop — works but serialises the computation onto one core, leaving a multi-core worker mostly idle while a single field takes minutes it should take seconds. When this task runs as one stage of a fleet-scale batch, that idle time multiplies across hundreds of fields.

Dask solves both problems at once. It chunks the raster into blocks, builds a lazy task graph over them, and schedules those blocks across every core of the worker under a hard memory ceiling — never holding more than a few chunks resident at a time. The result is bounded-memory, multi-core NDVI on an orthomosaic that does not fit in RAM. This is the fine-grained per-tile parallelism referenced in batch and async orthomosaic processing; it slots inside a single compute task of the queue built in orchestrating multi-flight jobs with Celery and Redis.

From chunk grid to task graph to workers An orthomosaic divided into a grid of chunks aligned with the file's internal tiles. Each chunk becomes an independent chain of read, compute index and write tasks in the graph, and the scheduler assigns those chains to four workers, so peak memory is set by the chunk size rather than by the size of the mosaic. Chunk grid, aligned to internal tiles chunk = 2048 × 2048 px · peak memory set here, not by the mosaic Task graph — one independent chain per chunk read window index math write block read window index math write block … one chain per chunk, no cross-chunk dependency 4 workers chunk 1 chunk 2 chunk 3 chunk 4 unaligned chunks force read amplification

Prerequisites

Only the Dask stack differs from the parent guide:

TEXT 
dask==2024.5.0
distributed==2024.5.0
rioxarray==0.15.5
xarray==2024.5.0
rasterio==1.3.10
numpy==1.26.4
BASH 
pip install dask==2024.5.0 distributed==2024.5.0 rioxarray==0.15.5 \
    xarray==2024.5.0 rasterio==1.3.10 numpy==1.26.4

Input requirements:

  • A multi-band orthomosaic GeoTIFF, ideally internally tiled (tiled=True, 256 or 512 block size) so chunk reads map to whole storage blocks.
  • An explicit projected CRS (e.g. EPSG:32615) carried in the file profile.
  • Known band positions for the sensor — MicaSense RedEdge-MX puts Red at band 3 and NIR at band 5.

Step-by-Step

Chunk size sets memory per worker and task count together Bars showing memory per chunk for four chunk sizes on a five-band float32 raster. Small chunks keep memory low but create so many tasks that scheduler overhead dominates; very large chunks reduce task count until workers exceed their memory limit and spill to disk. 512 px chunks many tiny tasks, scheduler overhead dominates 2048 px chunks 34 MB per chunk — the working range 4096 px chunks 134 MB per chunk, fewer tasks 8192 px chunks 537 MB — workers start spilling to disk Megabytes per chunk, five bands at float32. Multiply by the worker count for the fleet's peak.

Step 1 — Start a bounded-memory LocalCluster

The LocalCluster is the parallelism boundary. Set memory_limit per worker so Dask spills to disk instead of the OS killing the process, and match n_workers/threads_per_worker to the machine. Inside a Celery worker you would typically use a single-worker, multi-thread cluster so the two schedulers do not oversubscribe cores.

PYTHON 
from dask.distributed import Client, LocalCluster

cluster = LocalCluster(
    n_workers=1,
    threads_per_worker=4,
    memory_limit="4GB",     # hard cap — spill to disk beyond this
    processes=False,
)
client = Client(cluster)
print(client.dashboard_link)

Step 2 — Open the orthomosaic lazily as a chunked array

rioxarray opens the raster as an xarray DataArray backed by Dask, preserving the CRS and transform. The chunks argument controls the block size; align it to the file’s internal tiling. Crucially, nothing is read yet — this builds the graph only.

PYTHON 
import rioxarray

CHUNK = 1024
da = rioxarray.open_rasterio(
    "ortho_field_north80.tif",
    chunks={"band": 1, "x": CHUNK, "y": CHUNK},
    lock=False,
)
# Lazy: no pixels loaded. CRS and shape are available from metadata.
assert da.rio.crs is not None, "Orthomosaic has no CRS"
print(f"Shape {da.shape}, chunks {da.chunks[1][:3]}..., CRS EPSG:{da.rio.crs.to_epsg()}")

Step 3 — Compute NDVI blockwise with map_blocks

Select the red and NIR bands (still lazy) and apply the ratio. Because dask.array overloads arithmetic, (nir - red) / (nir + red) builds a graph that runs per block. Converting nodata to NaN and guarding the division inside a map_blocks function keeps edge chunks correct.

PYTHON 
import numpy as np
import dask.array as dskarr

RED_BAND, NIR_BAND = 3, 5   # 1-indexed MicaSense RedEdge-MX
NODATA = da.rio.nodata

red = da.sel(band=RED_BAND).data.astype("float32")   # dask array, lazy
nir = da.sel(band=NIR_BAND).data.astype("float32")


def _ndvi_block(nir_b, red_b, nodata):
    nir_b = nir_b.copy()
    red_b = red_b.copy()
    if nodata is not None:
        nir_b[nir_b == nodata] = np.nan
        red_b[red_b == nodata] = np.nan
    with np.errstate(invalid="ignore", divide="ignore"):
        out = (nir_b - red_b) / (nir_b + red_b)
    return out.astype("float32")


ndvi = dskarr.map_blocks(_ndvi_block, nir, red, NODATA, dtype="float32")
# Still lazy — ndvi is a graph, not an array of pixels
print(f"NDVI graph: {ndvi.shape}, {ndvi.numblocks} blocks")

Step 4 — Write and verify (complete runnable script)

Wrapping the lazy NDVI back into a georeferenced DataArray lets rioxarray write it with the correct CRS and transform. The write triggers the single compute that streams blocks to disk under the memory cap.

PYTHON 
import numpy as np
import xarray as xr
import rioxarray
import dask.array as dskarr
from dask.distributed import Client, LocalCluster

RED_BAND, NIR_BAND = 3, 5
CHUNK = 1024
SRC = "ortho_field_north80.tif"
DST = "ndvi_field_north80.tif"


def _ndvi_block(nir_b, red_b, nodata):
    nir_b, red_b = nir_b.copy(), red_b.copy()
    if nodata is not None:
        nir_b[nir_b == nodata] = np.nan
        red_b[red_b == nodata] = np.nan
    with np.errstate(invalid="ignore", divide="ignore"):
        return ((nir_b - red_b) / (nir_b + red_b)).astype("float32")


if __name__ == "__main__":
    cluster = LocalCluster(n_workers=1, threads_per_worker=4,
                           memory_limit="4GB", processes=False)
    client = Client(cluster)

    da = rioxarray.open_rasterio(
        SRC, chunks={"band": 1, "x": CHUNK, "y": CHUNK}, lock=False,
    )
    src_epsg = da.rio.crs.to_epsg()
    assert src_epsg is not None, "Source orthomosaic has no CRS"

    red = da.sel(band=RED_BAND).data.astype("float32")
    nir = da.sel(band=NIR_BAND).data.astype("float32")
    nodata = da.rio.nodata

    ndvi = dskarr.map_blocks(_ndvi_block, nir, red, nodata, dtype="float32")

    # Rebuild a georeferenced DataArray for output
    ndvi_da = xr.DataArray(
        ndvi,
        coords={"y": da["y"], "x": da["x"]},
        dims=("y", "x"),
    )
    ndvi_da.rio.write_crs(f"EPSG:{src_epsg}", inplace=True)
    ndvi_da.rio.write_nodata(np.nan, inplace=True)

    # The write is the single compute — blocks stream to disk, memory stays bounded
    ndvi_da.rio.to_raster(DST, tiled=True, blockxsize=256, blockysize=256,
                          compress="deflate", dtype="float32")

    # ── Verification: shape and CRS must survive the round trip ──
    with rioxarray.open_rasterio(DST) as out:
        assert out.rio.crs.to_epsg() == src_epsg, (
            f"CRS drift: in EPSG:{src_epsg}, out EPSG:{out.rio.crs.to_epsg()}"
        )
        assert out.shape[-2:] == da.shape[-2:], (
            f"Shape mismatch: in {da.shape[-2:]}, out {out.shape[-2:]}"
        )
        print(f"OK — NDVI {out.shape[-2:]} EPSG:{out.rio.crs.to_epsg()} written to {DST}")

    client.close()
    cluster.close()

Gotchas & Edge Cases

  • Triggering compute too early. Calling .compute(), .values, or np.asarray() on the lazy array before the final write materialises the whole raster in RAM — the exact MemoryError Dask exists to prevent. Keep everything lazy until to_raster.
Why chunk boundaries should follow the file's own tiling Two panels comparing Dask chunks aligned to a raster's internal tile grid with unaligned chunks. Aligned chunks read exactly the blocks they need; unaligned chunks pull partial blocks from neighbours, so border blocks are read several times and total I/O rises with no visible change in the code. Chunks aligned to the file's tiles Chunk boundaries fall on internal tile edges Each chunk reads exactly the blocks it needs. No block is read twice. I/O is proportional to the data. Chunks unaligned Chunk boundaries cut across internal tiles Each chunk pulls partial blocks from neighbours. Border blocks are read two or four times. Read amplification, invisible in the code.

  • Chunk size misaligned with the GeoTIFF. If the Dask chunk (e.g. 1000) does not tile evenly over the file’s internal blocks (e.g. 256), each chunk read spans partial blocks and reads them repeatedly — read amplification that can triple I/O. Use a chunk that is a multiple of the internal block size, or read the block size from src.block_shapes and match it.

  • NaN edges from unconverted nodata. Partial edge chunks contain nodata that, if left as its raw integer sentinel, divides into spurious NDVI values. Convert nodata to NaN inside the mapped function and set the output nodata to NaN, as the script does, so downstream cloud masking and index thresholds treat the edges correctly.

  • Oversubscribing cores inside a Celery worker. A prefork Celery worker already pins a core per process; adding a multi-process Dask LocalCluster on top oversubscribes the CPU and thrashes. Use processes=False with a small threads_per_worker so Dask threads share the single Celery worker’s core budget.

Parent Guide

This guide is part of Batch and Async Orthomosaic Processing — see there for the full architecture, the job manifest, and how this per-tile compute fits inside the distributed queue.

Parent section Batch and Async Orthomosaic Processing Related Orchestrating Multi-Flight Jobs with Celery and Redis